import numpy as np
import pyqmc.gpu as gpu
"""
Collection of 3d function objects. Each has a dictionary parameters, which corresponds
to any variational parameters the funtion has.
They should implement the following functions, all of which take input values (x, r).
x should be of dimension (nconf,...,3).
value(x):
returns f(x)
gradient(x)
returns grad f(x) (nconf,...,3)
laplacian(x)
returns diagonals of Hessian (nconf,...,3)
pgradient(x)
returns dp f(x) as a dictionary corresponding to the keys of self.parameters
"""
[docs]class GaussianFunction:
r"""A representation of a Gaussian:
:math:`\exp(-\alpha r^2)`
where :math:`\alpha` can be accessed through parameters['exponent']
"""
def __init__(self, exponent):
self.parameters = {"exponent": exponent}
[docs] def value(self, x, r):
"""Returns function exp(-exponent*r^2).
:parameter x: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: function value
:rtype: (nconfig,...) array
"""
return np.exp(-self.parameters["exponent"] * r * r)
[docs] def gradient(self, x, r):
"""Returns gradient of function.
:parameter x: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient
:rtype: (nconfig,...,3)
"""
v = self.value(x, r)
return -2 * self.parameters["exponent"] * x * v[..., np.newaxis]
[docs] def gradient_value(self, x, r):
"""
:parameter x: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and value
:rtype: tuple of (nconfig,...,3) arrays
"""
v = self.value(x, r)
g = -2 * self.parameters["exponent"] * x * v[..., np.newaxis]
return g, v
[docs] def laplacian(self, x, r):
"""Returns laplacian of function.
:parameter x: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: laplacian (components of laplacian d^2/dx_i^2 separately)
:rtype: (nconfig,...,3)
"""
v = self.value(x, r)
alpha = self.parameters["exponent"]
return (4 * alpha * alpha * x * x - 2 * alpha) * v[..., np.newaxis]
[docs] def gradient_laplacian(self, x, r):
"""Returns gradient and laplacian of function.
:parameter x: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and laplacian
:rtype: tuple of two (nconfig,...,3) arrays (components of laplacian d^2/dx_i^2 separately)
"""
v = self.value(x, r)[..., np.newaxis]
alpha = self.parameters["exponent"]
grad = -2 * alpha * x * v
lap = (4 * alpha * alpha * x * x - 2 * alpha) * v
return grad, lap
[docs] def pgradient(self, x, r):
"""Returns parameters gradient.
:parameter x: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: parameter gradient {'exponent': d/dexponent}
:rtype: dictionary
"""
r2 = r * r
return {"exponent": -r2 * np.exp(-self.parameters["exponent"] * r2)}
[docs]class PadeFunction:
r"""
a_k(r) = (alpha_k*r/(1+alpha_k*r))^2
alpha_k = alpha/2^k, k starting at 0
:math:`a_k(r) = \left( \frac{\alpha_k r}{1 + \alpha_k r} \right)^2`
where
:math:`\alpha_k = \frac{\alpha}{2^k}`, :math:`k` starting at 0
"""
def __init__(self, alphak):
self.parameters = {"alphak": alphak}
[docs] def value(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: function value
:rtype: (nconfig,...) array
"""
a = self.parameters["alphak"] * r
return (a / (1 + a)) ** 2
[docs] def gradient(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient
:rtype: (nconfig,...,3) array
"""
a = self.parameters["alphak"] * r[..., np.newaxis]
return 2 * self.parameters["alphak"] ** 2 / (1 + a) ** 3 * rvec
[docs] def gradient_value(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and value
:rtype: tuple of (nconfig,...,3) arrays
"""
a = self.parameters["alphak"] * r
value = (a / (1 + a)) ** 2
a = a[..., np.newaxis]
grad = 2 * self.parameters["alphak"] ** 2 / (1 + a) ** 3 * rvec
return grad, value
[docs] def laplacian(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: laplacian (returns components of laplacian d^2/dx_i^2 separately)
:rtype: (nconfig,...,3) array
"""
a = self.parameters["alphak"] * r[..., np.newaxis]
# lap = 6*self.parameters['alphak']**2 * (1+a)**(-4) #scalar formula
lap = (
2
* self.parameters["alphak"] ** 2
* (1 + a) ** (-3)
* (1 - 3 * a / (1 + a) * (rvec / r[..., np.newaxis]) ** 2)
)
return lap
[docs] def gradient_laplacian(self, rvec, r):
"""Returns gradient and laplacian of function.
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and laplacian (returns components of laplacian d^2/dx_i^2 separately)
:rtype: tuple of (nconfig,...,3) arrays
"""
a = self.parameters["alphak"] * r[..., np.newaxis]
temp = 2 * self.parameters["alphak"] ** 2 / (1 + a) ** 3
grad = temp * rvec
lap = temp * (1 - 3 * a / (1 + a) * (rvec / r[..., np.newaxis]) ** 2)
return grad, lap
[docs] def pgradient(self, rvec, r):
"""Returns gradient with respect to parameter alphak
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: parameter gradient {'alphak': akderiv}
:rtype: dictionary
"""
a = self.parameters["alphak"] * r
akderiv = 2 * a / (1 + a) ** 3 * r
return {"alphak": akderiv}
@gpu.fuse()
def polypadevalue(z, beta):
p = z * z * (6 - 8 * z + 3 * z * z)
return (1 - p) / (1 + beta * p)
@gpu.fuse()
def polypadegradvalue(r, beta, rcut):
z = r / rcut
p = z * z * (6 - 8 * z + 3 * z * z)
dpdz = 12 * z * (z * z - 2 * z + 1)
dbdp = -(1 + beta) / (1 + beta * p) ** 2
dzdx_rvec = 1 / (r * rcut)
grad_rvec = dbdp * dpdz * dzdx_rvec
value = (1 - p) / (1 + beta * p)
return grad_rvec, value
[docs]class PolyPadeFunction:
r"""
:math:`b(r) = \frac{1-p(z)}{1+\beta p(z)}`
:math:`z = r/r_{\rm cut}`
where
:math:`p(z) = 6z^2 - 8z^3 + 3z^4`
This function is positive at small r, and is zero for :math:`r \ge r_{\rm cut}`.
"""
def __init__(self, beta, rcut):
self.parameters = {
"beta": gpu.cp.asarray(beta),
"rcut": gpu.cp.asarray(rcut),
}
[docs] def value(self, rvec, r):
"""Returns function (1-p(r/rcut))/(1+beta*p(r/rcut))
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: function value
:rtype: (nconfig,...) array
"""
mask = r < self.parameters["rcut"]
z = r[mask] / self.parameters["rcut"]
func = gpu.cp.zeros(r.shape)
func[mask] = polypadevalue(z, self.parameters["beta"])
return func
[docs] def gradient_value(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and value
:rtype: tuple of (nconfig,...,3) arrays
"""
value = gpu.cp.zeros(r.shape)
grad = gpu.cp.zeros(rvec.shape)
mask = r < self.parameters["rcut"]
grad_rvec, value[mask] = polypadegradvalue(
r[mask],
self.parameters["beta"],
self.parameters["rcut"],
)
grad[mask] = np.einsum("ij,i->ij", rvec[mask], grad_rvec)
return grad, value
[docs] def gradient(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient
:rtype: (nconfig,...,3) array
"""
grad = gpu.cp.zeros(rvec.shape)
mask = r < self.parameters["rcut"]
r = r[mask][..., np.newaxis]
rvec = rvec[mask]
z = r / self.parameters["rcut"]
p = z * z * (6 - 8 * z + 3 * z * z)
dpdz = 12 * z * (z * z - 2 * z + 1)
dbdp = -(1 + self.parameters["beta"]) / (1 + self.parameters["beta"] * p) ** 2
dzdx = rvec / (r * self.parameters["rcut"])
grad[mask] = dbdp * dpdz * dzdx
return grad
[docs] def laplacian(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: laplacian
returns components of laplacian d^2/dx_i^2 separately
:rtype: (nconfig,...,3) array
"""
return self.gradient_laplacian(rvec, r)[1]
[docs] def gradient_laplacian(self, rvec, r):
"""Returns gradient and laplacian of function.
:parameter x: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and laplacian
:rtype: tuple of two (nconfig,...,3) arrays (components of laplacian d^2/dx_i^2 separately)
"""
grad = gpu.cp.zeros(rvec.shape)
lap = gpu.cp.zeros(rvec.shape)
mask = r < self.parameters["rcut"]
r = r[..., np.newaxis]
r = r[mask]
rvec = rvec[mask]
z = r / self.parameters["rcut"]
beta = self.parameters["beta"]
p = z * z * (6 - 8 * z + 3 * z * z)
dpdz = 12 * z * (z * z - 2 * z + 1)
dbdp = -(1 + beta) / (1 + beta * p) ** 2
dzdx = rvec / (r * self.parameters["rcut"])
gradmask = dbdp * dpdz * dzdx
d2pdz2_over_dpdz = (3 * z - 1) / (z * (z - 1))
d2bdp2_over_dbdp = -2 * beta / (1 + beta * p)
d2zdx2 = (1 - (rvec / r) ** 2) / (r * self.parameters["rcut"])
grad[mask] = gradmask
lap[mask] += dbdp * dpdz * d2zdx2 + (
gradmask * (d2bdp2_over_dbdp * dpdz * dzdx + d2pdz2_over_dpdz * dzdx)
)
return grad, lap
[docs] def pgradient(self, rvec, r):
"""Returns gradient of self.value with respect to all parameters
:parameter rvec: (nconf,...,3)
:parameter r: (nconf,...)
:return paramderivs: dictionary {'rcut':d/drcut,'beta':d/dbeta}
"""
mask = r >= self.parameters["rcut"]
z = r / self.parameters["rcut"]
beta = self.parameters["beta"]
p = z * z * (6 - 8 * z + 3 * z * z)
dbdp = -(1 + beta) / (1 + beta * p) ** 2
dpdz = 12 * z * (z * z - 2 * z + 1)
derivrcut = dbdp * dpdz * (-z / self.parameters["rcut"])
derivbeta = -p * (1 - p) / (1 + beta * p) ** 2
derivrcut[mask] = 0.0
derivbeta[mask] = 0.0
pderiv = {"rcut": derivrcut, "beta": derivbeta}
return pderiv
[docs]class CutoffCuspFunction:
r"""
:math:`b(r) = -\frac{p(r/r_{\rm cut})}{1+\gamma*p(r/r_{\rm cut})} + \frac{1}{3+\gamma}`
where
:math:`p(y) = y - y^2 + y^3/3`
This function is positive at small r, and is zero for :math:`r \ge r_{\rm cut}`.
"""
def __init__(self, gamma, rcut):
self.parameters = {"gamma": gamma, "rcut": rcut}
[docs] def value(self, rvec, r):
"""Returns function p(r/rcut)/(1+gamma*p(r/rcut))
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: function value
:rtype: (nconfig,...) array
"""
mask = r < self.parameters["rcut"]
y = r[mask] / self.parameters["rcut"]
gamma = self.parameters["gamma"]
p = y - y * y + y * y * y / 3
func = gpu.cp.zeros(r.shape)
func[mask] = -p / (1 + gamma * p) + 1 / (3 + gamma)
return func * self.parameters["rcut"]
[docs] def gradient(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient
:rtype: (nconfig,...,3) array
"""
rcut = self.parameters["rcut"]
gamma = self.parameters["gamma"]
mask = r < rcut
r = r[mask][..., np.newaxis]
y = r / rcut
a = 1 - 2 * y + y * y
b = y - y * y + y * y * y / 3
c = 1 / (1 + gamma * b) ** 2 / (rcut * r)
grad = gpu.cp.zeros(rvec.shape)
grad[mask] = -rvec[mask] * a * c * rcut
return grad
[docs] def gradient_value(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and value
:rtype: tuple of (nconfig,...,3) arrays
"""
grad = gpu.cp.zeros(rvec.shape)
value = gpu.cp.zeros(r.shape)
rcut = self.parameters["rcut"]
gamma = self.parameters["gamma"]
mask = r < rcut
r = r[mask][..., np.newaxis]
y = r / rcut
a = 1 - 2 * y + y * y
b = y - y * y + y * y * y / 3
c = 1 / (1 + gamma * b) ** 2 / (rcut * r)
grad[mask] = -rvec[mask] * a * c * rcut
b = b[..., 0]
value[mask] = -b / (1 + gamma * b) + 1 / (3 + gamma)
return grad, value * rcut
[docs] def laplacian(self, rvec, r):
"""
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: laplacian (returns components of laplacian d^2/dx_i^2 separately)
:rtype: (nconfig,...,3) array
"""
lap = gpu.cp.zeros(rvec.shape)
rcut = self.parameters["rcut"]
gamma = self.parameters["gamma"]
mask = r < rcut
r = r[mask][..., np.newaxis]
rvec = rvec[mask]
y = r / rcut
a = 1 - 2 * y + y * y
b = y - y * y + y * y * y / 3
c = 1 / (1 + gamma * b) ** 2 / (rcut * r)
temp = 2 * (y - 1) / (rcut * r)
temp -= a / r ** 2
temp -= 2 * a * a * c * gamma * (1 + gamma * b)
lap[mask] = -rcut * c * (a + rvec ** 2 * temp)
return lap
[docs] def gradient_laplacian(self, rvec, r):
"""Returns gradient and laplacian of function.
:parameter rvec: (nconfig,...,3)
:parameter r: (nconfig,...)
:returns: gradient and laplacian
:rtype: tuple of two (nconfig,...,3) arrays (components of laplacian d^2/dx_i^2 separately)
"""
grad = gpu.cp.zeros(rvec.shape)
lap = gpu.cp.zeros(rvec.shape)
rcut = self.parameters["rcut"]
gamma = self.parameters["gamma"]
mask = r < rcut
r = r[mask][..., np.newaxis]
rvec = rvec[mask]
y = r / rcut
a = 1 - 2 * y + y * y
b = y - y * y + y * y * y / 3
c = 1 / (1 + gamma * b) ** 2 / (rcut * r)
grad[mask] = -rcut * a * c * rvec
temp = 2 * (y - 1) / (rcut * r)
temp -= a / r ** 2
temp -= 2 * a * a * c * gamma * (1 + gamma * b)
lap[mask] = -rcut * c * (a + rvec ** 2 * temp)
return grad, lap
[docs] def pgradient(self, rvec, r):
"""Returns gradient of self.value with respect to all parameters
:parameter rvec: (nconf,...,3)
:parameter r: (nconfig,...)
:returns: parameter derivatives {'rcut':d/drcut,'gamma':d/dgamma}
:rtype: dict
"""
rcut = self.parameters["rcut"]
gamma = self.parameters["gamma"]
mask = r > rcut
r = r
y = r / rcut
a = 1 - 2 * y + y * y
b = y - y * y + y * y * y / 3
c = a / (1 + gamma * b) ** 2 / (rcut * r)
val = -b / (1 + gamma * b) + 1 / (3 + gamma)
dfdrcut = y * a / (1 + gamma * b) ** 2 + val
dfdgamma = ((b / (1 + gamma * b)) ** 2 - 1 / (3 + gamma) ** 2) * rcut
dfdrcut[mask] = 0.0
dfdgamma[mask] = 0.0
func = {"rcut": dfdrcut, "gamma": dfdgamma}
return func
def test_func3d_gradient(bf, delta=1e-5):
rvec = gpu.cp.asarray(np.random.randn(150, 5, 10, 3)) # Internal indices irrelevant
r = np.linalg.norm(rvec, axis=-1)
grad = bf.gradient(rvec, r)
numeric = gpu.cp.zeros(rvec.shape)
for d in range(3):
pos = rvec.copy()
pos[..., d] += delta
plusval = bf.value(pos, np.linalg.norm(pos, axis=-1))
pos[..., d] -= 2 * delta
minuval = bf.value(pos, np.linalg.norm(pos, axis=-1))
numeric[..., d] = (plusval - minuval) / (2 * delta)
maxerror = np.max(np.abs(grad - numeric))
return gpu.asnumpy(maxerror)
def test_func3d_laplacian(bf, delta=1e-5):
rvec = gpu.cp.asarray(np.random.randn(150, 5, 10, 3)) # Internal indices irrelevant
r = np.linalg.norm(rvec, axis=-1)
lap = bf.laplacian(rvec, r)
numeric = gpu.cp.zeros(rvec.shape)
for d in range(3):
pos = rvec.copy()
pos[..., d] += delta
r = np.linalg.norm(pos, axis=-1)
plusval = bf.gradient(pos, r)[..., d]
pos[..., d] -= 2 * delta
r = np.linalg.norm(pos, axis=-1)
minuval = bf.gradient(pos, r)[..., d]
numeric[..., d] = (plusval - minuval) / (2 * delta)
maxerror = np.max(np.abs(lap - numeric))
return gpu.asnumpy(maxerror)
def test_func3d_gradient_laplacian(bf):
rvec = gpu.cp.asarray(np.random.randn(150, 10, 3))
r = np.linalg.norm(rvec, axis=-1)
grad = bf.gradient(rvec, r)
lap = bf.laplacian(rvec, r)
andgrad, andlap = bf.gradient_laplacian(rvec, r)
graderr = np.amax(np.abs(grad - andgrad))
laperr = np.amax(np.abs(lap - andlap))
return {"grad": graderr, "lap": laperr}
def test_func3d_gradient_value(bf):
rvec = gpu.cp.asarray(np.random.randn(150, 10, 3))
r = np.linalg.norm(rvec, axis=-1)
grad = bf.gradient(rvec, r)
val = bf.value(rvec, r)
andgrad, andval = bf.gradient_value(rvec, r)
graderr = np.linalg.norm((grad - andgrad))
valerr = np.linalg.norm((val - andval))
return {"grad": graderr, "val": valerr}
def test_func3d_pgradient(bf, delta=1e-5):
rvec = gpu.cp.asarray(np.random.randn(150, 10, 3))
r = np.linalg.norm(rvec, axis=-1)
pgrad = bf.pgradient(rvec, r)
numeric = {k: gpu.cp.zeros(v.shape) for k, v in pgrad.items()}
maxerror = {k: np.zeros(v.shape) for k, v in pgrad.items()}
for k in pgrad.keys():
bf.parameters[k] += delta
plusval = bf.value(rvec, r)
bf.parameters[k] -= 2 * delta
minuval = bf.value(rvec, r)
bf.parameters[k] += delta
numeric[k] = (plusval - minuval) / (2 * delta)
maxerror[k] = gpu.asnumpy(np.max(np.abs(pgrad[k] - numeric[k])))
if maxerror[k] > 1e-5:
print(k, "\n", pgrad[k] - numeric[k])
return maxerror