import numpy as np
import pyqmc.gpu as gpu
[docs]class JastrowSpin:
r"""
1 body and 2 body jastrow factor
The Jastrow form is :math:`e^{U(R)}`, where
.. math:: U(R) = \sum_{I, i, k} c^{a}_{Ik\sigma(i)} a_{k}(r_{Ii}) + \sum_{i,j,k} c^{b}_{k\sigma(i)\sigma(j)} b^{l}(r_{ij})
"""
def __init__(self, mol, a_basis, b_basis):
r"""
Args:
mol : a pyscf molecule object
a_basis : list of func3d objects that comprise the electron-ion basis
b_basis : list of func3d objects that comprise the electron-electron basis
"""
self.a_basis = a_basis
self.b_basis = b_basis
self.parameters = {}
self._nelec = np.sum(mol.nelec)
self._mol = mol
self.parameters["bcoeff"] = gpu.cp.zeros((len(b_basis), 3))
self.parameters["acoeff"] = gpu.cp.zeros((self._mol.natm, len(a_basis), 2))
self.iscomplex = False
[docs] def recompute(self, configs):
r"""
_avalues is the array for current configurations :math:`A_{Iks} = \sum_s a_{k}(r_{Is})` where :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums.
_bvalues is the array for current configurations :math:`B_{ls} = \sum_s b_{l}(r_{s})` where :math:`s` indexes over :math:`\uparrow\uparrow` (:math:`\alpha_1 < \alpha_2`), :math:`\uparrow\downarrow` (:math:`\alpha, \beta`), and :math:`\downarrow\downarrow` (:math:`\beta_1 < \beta_2`) sums.
the partial sums store values before summing over electrons
_a_partial is the array :math:`A^p_{eIk} = a_k(r_{Ie}`, where :math:`e` is any electron
_b_partial is the array :math:`B^p_{els} = \sum_s b_l(r_{es}`, where :math:`e` is any electron, :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums, not including :math:`e`.
"""
self._configscurrent = configs.copy()
nconf, nelec = configs.configs.shape[:2]
nexpand = len(self.b_basis)
aexpand = len(self.a_basis)
self._bvalues = gpu.cp.zeros((nconf, nexpand, 3))
self._avalues = gpu.cp.zeros((nconf, self._mol.natm, aexpand, 2))
self._a_partial = gpu.cp.zeros((nelec, nconf, self._mol.natm, aexpand))
self._b_partial = gpu.cp.zeros((nelec, nconf, nexpand, 2))
notmask = np.ones(nconf, dtype=bool)
for e in range(nelec):
epos = configs.electron(e)
self._a_partial[e] = self._a_update(e, epos, notmask)
self._b_partial[e] = self._b_update(e, epos, notmask)
# electron-electron distances
nup = self._mol.nelec[0]
d_upup, ij = configs.dist.dist_matrix(configs.configs[:, :nup])
d_updown, ij = configs.dist.pairwise(
configs.configs[:, :nup], configs.configs[:, nup:]
)
d_downdown, ij = configs.dist.dist_matrix(configs.configs[:, nup:])
# Update bvalues according to spin case
for j, d in enumerate([d_upup, d_updown, d_downdown]):
d = gpu.cp.asarray(d)
r = gpu.cp.linalg.norm(d, axis=-1)
for i, b in enumerate(self.b_basis):
self._bvalues[:, i, j] = gpu.cp.sum(b.value(d, r), axis=1)
# electron-ion distances
di = gpu.cp.zeros((nelec, nconf, self._mol.natm, 3))
for e in range(nelec):
di[e] = gpu.cp.asarray(
configs.dist.dist_i(self._mol.atom_coords(), configs.configs[:, e, :])
)
ri = gpu.cp.linalg.norm(di, axis=-1)
# Update avalues according to spin case
for i, a in enumerate(self.a_basis):
avals = a.value(di, ri)
self._avalues[:, :, i, 0] = gpu.cp.sum(avals[:nup], axis=0)
self._avalues[:, :, i, 1] = gpu.cp.sum(avals[nup:], axis=0)
u = gpu.cp.sum(self._bvalues * self.parameters["bcoeff"], axis=(2, 1))
u += gpu.cp.einsum("ijkl,jkl->i", self._avalues, self.parameters["acoeff"])
return (np.ones(len(u)), gpu.asnumpy(u))
[docs] def updateinternals(self, e, epos, wrap=None, mask=None):
r"""Update a and b sums.
_avalues is the array for current configurations :math:`A_{Iks} = \sum_s a_{k}(r_{Is})` where :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums.
_bvalues is the array for current configurations :math:`B_{ls} = \sum_s b_{l}(r_{s})` where :math:`s` indexes over :math:`\uparrow\uparrow` (:math:`\alpha_1 < \alpha_2`), :math:`\uparrow\downarrow` (:math:`\alpha, \beta`), and :math:`\downarrow\downarrow` (:math:`\beta_1 < \beta_2`) sums.
The update for _avalues and _b_values from moving one electron only requires computing the new sum for that electron. The sums for the electron in the current configuration are stored in _a_partial and _b_partial"""
if mask is None:
mask = [True] * self._configscurrent.configs.shape[0]
edown = int(e >= self._mol.nelec[0])
aupdate = self._a_update(e, epos, mask)
bupdate = self._b_update(e, epos, mask)
self._avalues[:, :, :, edown][mask] += aupdate - self._a_partial[e][mask]
self._bvalues[:, :, edown : edown + 2][mask] += (
bupdate - self._b_partial[e][mask]
)
self._a_partial[e][mask] = aupdate
self._update_b_partial(e, epos, mask)
self._configscurrent.move(e, epos, mask)
def _a_update(self, e, epos, mask):
r"""
Calculate a (e-ion) partial sum for electron e
_a_partial_e is the array :math:`A^p_{iIk} = a_k(r^i_{Ie}` with e fixed
i is the configuration index
Args:
e: fixed electron index
epos: configs object for electron e
mask: mask over configs axis, only return values for configs where mask==True. a_partial_e might have a smaller configs axis than epos, _configscurrent, and _a_partial because of the mask.
"""
d = gpu.cp.asarray(
epos.dist.dist_i(self._mol.atom_coords(), epos.configs[mask])
)
r = gpu.cp.linalg.norm(d, axis=-1)
a_partial_e = gpu.cp.zeros((*r.shape, self._a_partial.shape[3]))
for k, a in enumerate(self.a_basis):
a_partial_e[..., k] = a.value(d, r)
return a_partial_e
def _b_update(self, e, epos, mask):
r"""
Calculate b (e-e) partial sums for electron e
_b_partial_e is the array :math:`B^p_{ils} = \sum_s b_l(r^i_{es}`, with e fixed; :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums, not including electron e.
:math:`i` is the configuration index.
Args:
e: fixed electron index
epos: configs object for electron e
mask: mask over configs axis, only return values for configs where mask==True. b_partial_e might have a smaller configs axis than epos, _configscurrent, and _b_partial because of the mask.
"""
nup = self._mol.nelec[0]
sep = nup - int(e < nup)
not_e = np.arange(self._nelec) != e
d = gpu.cp.asarray(
epos.dist.dist_i(
self._configscurrent.configs[mask][:, not_e], epos.configs[mask]
)
)
r = gpu.cp.linalg.norm(d, axis=-1)
b_partial_e = gpu.cp.zeros((*r.shape[:-1], *self._b_partial.shape[2:]))
for l, b in enumerate(self.b_basis):
bval = b.value(d, r)
b_partial_e[..., l, 0] = bval[..., :sep].sum(axis=-1)
b_partial_e[..., l, 1] = bval[..., sep:].sum(axis=-1)
return b_partial_e
def _b_update_many(self, e, epos, mask, spin):
r"""
Compute the update to b for each electron moving to epos.
Calculate b (e-e) partial sums for electron e
_b_partial_e is the array :math:`B^p_{ils} = \sum_s b_l(r^i_{es}`, with e fixed; :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums, not including electron e.
:math:`i` is the configuration index.
Args:
e: fixed electron index
epos: configs object for electron e
mask: mask over configs axis, only return values for configs where mask==True. b_partial_e might have a smaller configs axis than epos, _configscurrent, and _b_partial because of the mask.
"""
nup = self._mol.nelec[0]
d = gpu.cp.asarray(
epos.dist.dist_i(self._configscurrent.configs[mask], epos.configs[mask])
)
r = gpu.cp.linalg.norm(d, axis=-1)
b_partial_e = gpu.cp.zeros(
(e.shape[0], *r.shape[:-1], *self._b_partial.shape[2:])
)
for l, b in enumerate(self.b_basis):
bval = b.value(d, r)
b_partial_e[..., l, 0] = bval[..., :nup].sum(axis=-1)
b_partial_e[..., l, 1] = bval[..., nup:].sum(axis=-1)
# b_partial_e[..., l, spin] -= bval[..., e].T
b_partial_e[..., l, spin] -= np.moveaxis(bval[..., e], -1, 0)
return b_partial_e
def _update_b_partial(self, e, epos, mask):
r"""
Calculate b (e-e) partial sum contributions from electron e
_b_partial_e is the array :math:`B^p_{ils} = \sum_s b_l(r^i_{es}`, with e fixed; :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums, not including electron e.
Since :math:`B^p_{ils}` is summed over other electrons, moving electron e will affect other partial sums. This function updates all the necessary partial sums instead of just evaluating the one for electron e.
:math:`i` is the configuration index.
Args:
e: fixed electron index
epos: configs object for electron e
mask: mask over configs axis, only return values for configs where mask==True. b_partial_e might have a smaller configs axis than epos, _configscurrent, and _b_partial because of the mask.
"""
nup = self._mol.nelec[0]
sep = nup - int(e < nup)
not_e = np.arange(self._nelec) != e
edown = int(e >= nup)
d = gpu.cp.asarray(
epos.dist.dist_i(
self._configscurrent.configs[mask][:, not_e], epos.configs[mask]
)
)
r = gpu.cp.linalg.norm(d, axis=-1)
dold = gpu.cp.asarray(
epos.dist.dist_i(
self._configscurrent.configs[mask][:, not_e],
self._configscurrent.configs[mask, e],
)
)
rold = gpu.cp.linalg.norm(dold, axis=-1)
b_partial_e = gpu.cp.zeros((np.sum(mask), *self._b_partial.shape[2:]))
eind, mind = np.ix_(not_e, mask)
for l, b in enumerate(self.b_basis):
bval = b.value(d, r)
bdiff = bval - b.value(dold, rold)
self._b_partial[eind, mind, l, edown] += bdiff.transpose((1, 0))
self._b_partial[e, :, l, 0][mask] = bval[:, :sep].sum(axis=1)
self._b_partial[e, :, l, 1][mask] = bval[:, sep:].sum(axis=1)
[docs] def value(self):
"""Compute the current log value of the wavefunction"""
u = gpu.cp.sum(self._bvalues * self.parameters["bcoeff"], axis=(2, 1))
u += gpu.cp.einsum("ijkl,jkl->i", self._avalues, self.parameters["acoeff"])
return (np.ones(len(u)), gpu.asnumpy(u))
[docs] def gradient(self, e, epos):
r"""We compute the gradient for electron e as
:math:`\nabla_e \ln \Psi_J = \sum_k c_k \left(\sum_{j > e} \nabla_e b_k(r_{ej}) + \sum_{i < e} \nabla_e b_k(r_{ie})\right)`
So we need to compute the gradient of the b's for these indices.
Note that we need to compute distances between electron position given and the current electron distances.
We will need this for laplacian() as well"""
nconf, nelec = self._configscurrent.configs.shape[:2]
nup = self._mol.nelec[0]
# Get e-e and e-ion distances
not_e = np.arange(nelec) != e
dnew = gpu.cp.asarray(
epos.dist.dist_i(self._configscurrent.configs, epos.configs)[:, not_e]
)
dinew = gpu.cp.asarray(epos.dist.dist_i(self._mol.atom_coords(), epos.configs))
rnew = gpu.cp.linalg.norm(dnew, axis=-1)
rinew = gpu.cp.linalg.norm(dinew, axis=-1)
grad = gpu.cp.zeros((3, nconf))
# Check if selected electron is spin up or down
eup = int(e < nup)
edown = int(e >= nup)
for c, b in zip(self.parameters["bcoeff"], self.b_basis):
bgrad = b.gradient(dnew, rnew)
grad += c[edown] * gpu.cp.sum(bgrad[:, : nup - eup], axis=1).T
grad += c[1 + edown] * gpu.cp.sum(bgrad[:, nup - eup :], axis=1).T
for c, a in zip(self.parameters["acoeff"].transpose()[edown], self.a_basis):
grad += gpu.cp.einsum("j,ijk->ki", c, a.gradient(dinew, rinew))
return gpu.asnumpy(grad)
def gradient_value(self, e, epos):
r""""""
nconf, nelec = self._configscurrent.configs.shape[:2]
nup = self._mol.nelec[0]
# Get e-e and e-ion distances
not_e = np.arange(nelec) != e
dnew = gpu.cp.asarray(
epos.dist.dist_i(self._configscurrent.configs[:, not_e], epos.configs)
)
dinew = gpu.cp.asarray(epos.dist.dist_i(self._mol.atom_coords(), epos.configs))
rnew = gpu.cp.linalg.norm(dnew, axis=-1)
rinew = gpu.cp.linalg.norm(dinew, axis=-1)
grad = gpu.cp.zeros((3, nconf))
# Check if selected electron is spin up or down
eup = int(e < nup)
edown = int(e >= nup)
b_partial_e = gpu.cp.zeros((*rnew.shape[:-1], *self._b_partial.shape[2:]))
for l, b in enumerate(self.b_basis):
c = self.parameters["bcoeff"][l]
bgrad, bval = b.gradient_value(dnew, rnew)
grad += c[edown] * gpu.cp.sum(bgrad[:, : nup - eup], axis=1).T
grad += c[1 + edown] * gpu.cp.sum(bgrad[:, nup - eup :], axis=1).T
b_partial_e[..., l, 0] = bval[..., : nup - eup].sum(axis=-1)
b_partial_e[..., l, 1] = bval[..., nup - eup :].sum(axis=-1)
a_partial_e = gpu.cp.zeros((*rinew.shape, self._a_partial.shape[3]))
for k, a in enumerate(self.a_basis):
c = self.parameters["acoeff"][:, k, edown]
agrad, aval = a.gradient_value(dinew, rinew)
grad += gpu.cp.einsum("j,ijk->ki", c, agrad)
a_partial_e[..., k] = aval
deltaa = a_partial_e - self._a_partial[e]
a_val = gpu.cp.einsum(
"...jk,jk->...", deltaa, self.parameters["acoeff"][..., edown]
)
deltab = b_partial_e - self._b_partial[e]
b_val = gpu.cp.einsum(
"...jk,jk->...", deltab, self.parameters["bcoeff"][:, edown : edown + 2]
)
val = gpu.cp.exp(b_val + a_val)
return gpu.asnumpy(grad), gpu.asnumpy(val)
[docs] def gradient_laplacian(self, e, epos):
""" """
nconf, nelec = self._configscurrent.configs.shape[:2]
nup = self._mol.nelec[0]
# Get e-e and e-ion distances
not_e = np.arange(nelec) != e
dnew = gpu.cp.asarray(
epos.dist.dist_i(self._configscurrent.configs, epos.configs)[:, not_e]
)
dinew = gpu.cp.asarray(epos.dist.dist_i(self._mol.atom_coords(), epos.configs))
rnew = gpu.cp.linalg.norm(dnew, axis=-1)
rinew = gpu.cp.linalg.norm(dinew, axis=-1)
eup = int(e < nup)
edown = int(e >= nup)
grad = gpu.cp.zeros((3, nconf))
lap = gpu.cp.zeros(nconf)
# a-value component
for c, a in zip(self.parameters["acoeff"].transpose()[edown], self.a_basis):
g, l = a.gradient_laplacian(dinew, rinew)
grad += gpu.cp.einsum("j,ijk->ki", c, g)
lap += gpu.cp.einsum("j,ijk->i", c, l)
# b-value component
for c, b in zip(self.parameters["bcoeff"], self.b_basis):
bgrad, blap = b.gradient_laplacian(dnew, rnew)
grad += c[edown] * gpu.cp.sum(bgrad[:, : nup - eup], axis=1).T
grad += c[1 + edown] * gpu.cp.sum(bgrad[:, nup - eup :], axis=1).T
lap += c[edown] * gpu.cp.sum(blap[:, : nup - eup], axis=(1, 2))
lap += c[1 + edown] * gpu.cp.sum(blap[:, nup - eup :], axis=(1, 2))
return gpu.asnumpy(grad), gpu.asnumpy(lap + gpu.cp.sum(grad ** 2, axis=0))
def laplacian(self, e, epos):
return self.gradient_laplacian(e, epos)[1]
[docs] def testvalue(self, e, epos, mask=None):
r"""
Compute the ratio :math:`\Psi_{\rm new}/\Psi_{\rm old}` for moving electron e to epos.
_avalues is the array for current configurations :math:`A_{Iks} = \sum_s a_{k}(r_{Is})` where :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums.
_bvalues is the array for current configurations :math:`B_{ls} = \sum_s b_{l}(r_{s})` where :math:`s` indexes over :math:`\uparrow\uparrow` (:math:`\alpha_1 < \alpha_2`), :math:`\uparrow\downarrow` (:math:`\alpha, \beta`), and :math:`\downarrow\downarrow` (:math:`\beta_1 < \beta_2`) sums.
The update for _avalues and _b_values from moving one electron only requires computing the new sum for that electron. The sums for the electron in the current configuration are stored in _a_partial and _b_partial.
deltaa = :math:`a_{k}(r_{Ie})`, indexing (atom, a_basis)
deltab = :math:`\sum_s b_{l}(r_{se})`, indexing (b_basis, spin s)
"""
if mask is None:
mask = [True] * epos.configs.shape[0]
edown = int(e >= self._mol.nelec[0])
deltaa = self._a_update(e, epos, mask) - self._a_partial[e][mask]
a_val = gpu.cp.einsum(
"...jk,jk->...", deltaa, self.parameters["acoeff"][..., edown]
)
deltab = self._b_update(e, epos, mask) - self._b_partial[e][mask]
b_val = gpu.cp.einsum(
"...jk,jk->...", deltab, self.parameters["bcoeff"][:, edown : edown + 2]
)
val = gpu.cp.exp(b_val + a_val)
if len(val.shape) == 2:
val = val.T
return gpu.asnumpy(val)
[docs] def testvalue_many(self, e, epos, mask=None):
r"""
Compute the ratio :math:`\Psi_{\rm new}/\Psi_{\rm old}` for moving electrons in e to epos.
_avalues is the array for current configurations :math:`A_{Iks} = \sum_s a_{k}(r_{Is})` where :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums.
_bvalues is the array for current configurations :math:`B_{ls} = \sum_s b_{l}(r_{s})` where :math:`s` indexes over :math:`\uparrow\uparrow` (:math:`\alpha_1 < \alpha_2`), :math:`\uparrow\downarrow` (:math:`\alpha, \beta`), and :math:`\downarrow\downarrow` (:math:`\beta_1 < \beta_2`) sums.
The update for _avalues and _b_values from moving one electron only requires computing the new sum for that electron. The sums for the electron in the current configuration are stored in _a_partial and _b_partial.
deltaa = :math:`a_{k}(r_{Ie})`, indexing (atom, a_basis)
deltab = :math:`\sum_s b_{l}(r_{se})`, indexing (b_basis, spin s)
"""
s = (e >= self._mol.nelec[0]).astype(int)
if mask is None:
mask = [True] * epos.configs.shape[0]
ratios = gpu.cp.zeros((epos.configs.shape[0], e.shape[0]))
for spin in [0, 1]:
ind = s == spin
deltaa = (
self._a_update(e[ind], epos, mask) - self._a_partial[e[ind]][:, mask]
)
deltab = (
self._b_update_many(e[ind], epos, mask, spin)
- self._b_partial[e[ind]][:, mask]
)
a_val = gpu.cp.einsum(
"...jk,jk->...", deltaa, self.parameters["acoeff"][..., spin]
)
b_val = gpu.cp.einsum(
"...jk,jk->...", deltab, self.parameters["bcoeff"][:, spin : spin + 2]
)
val = gpu.cp.exp(b_val + a_val)
if len(val.shape) == 2:
val = val.T
ratios[:, ind] = val
return gpu.asnumpy(ratios)
[docs] def pgradient(self):
"""Given the b sums, this is pretty trivial for the coefficient derivatives.
For the derivatives of basis functions, we will have to compute the derivative
of all the b's and redo the sums, similar to recompute()"""
return {
"bcoeff": gpu.asnumpy(self._bvalues),
"acoeff": gpu.asnumpy(self._avalues),
}
[docs] def u_components(self, rvec, r):
"""Given positions rvec and their magnitudes r, returns
dictionaries of the one-body and two-body Jastrow components.
Dictionaries are the spin components of U summed across the basis;
one-body also returns U for different atoms."""
u_onebody = {"up": [], "dn": []}
rvec = gpu.cp.asarray(rvec)
r = gpu.cp.asarray(r)
a_value = gpu.cp.asarray(list(map(lambda x: x.value(rvec, r), self.a_basis)))
u_onebody["up"] = gpu.cp.einsum(
"ij,jl->il", self.parameters["acoeff"][:, :, 0], a_value
)
u_onebody["dn"] = gpu.cp.einsum(
"ij,jl->il", self.parameters["acoeff"][:, :, 1], a_value
)
u_twobody = {"upup": [], "updn": [], "dndn": []}
b_value = gpu.cp.asarray(
list(map(lambda x: x.value(rvec, r), self.b_basis[1:]))
)
u_twobody["upup"] = gpu.cp.dot(self.parameters["bcoeff"][1:, 0], b_value)
u_twobody["updn"] = gpu.cp.dot(self.parameters["bcoeff"][1:, 1], b_value)
u_twobody["dndn"] = gpu.cp.dot(self.parameters["bcoeff"][1:, 2], b_value)
return u_onebody, u_twobody