def energy(self, configs):
r"""
Compute Coulomb energy for a set of configs.
.. math:: E_{\rm Coulomb} &= E_{\rm real+reciprocal}^{ee}
+ E_{\rm self+charged}^{ee}
\\&+ E_{\rm real+reciprocal}^{e\text{-ion}}
+ E_{\rm self+charged}^{e\text{-ion}}
\\&+ E_{\rm real+reciprocal}^{\text{ion-ion}}
+ E_{\rm self+charged}^{\text{ion-ion}}
:parameter configs: electron positions (walkers)
:type configs: (nconf, nelec, 3) PeriodicConfigs object
:returns:
* ee: electron-electron part
* ei: electron-ion part
* ii: ion-ion part
:rtype: float, float, float
"""
nelec = configs.configs.shape[1]
ee, ei = self.ewald_electron(configs)
ee += self.ee_const(nelec)
ei += self.ei_const(nelec)
ii = self.ion_ion + self.ii_const
return gpu.asnumpy(ee), gpu.asnumpy(ei), gpu.asnumpy(ii)
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 gradient_laplacian(self, e, epos):
_, e_grad, e_lap = self._get_val_grad_lap(epos)
lap1 = np.einsum(
"mn, dcm, cjn-> c",
self.parameters["gcoeff"],
e_lap[:, :, 0, :],
self.ao_val[:, :e, :],
optimize=self.optimize,
)
lap2 = np.einsum(
"mn, cim, dcn -> c",
self.parameters["gcoeff"],
self.ao_val[:, e + 1 :, :],
e_lap[:, :, 0, :],
optimize=self.optimize,
)
grad1 = np.einsum(
"mn, dcm, cjn -> dc",
self.parameters["gcoeff"],
e_grad[:, :, 0, :],
self.ao_val[:, :e, :],
optimize=self.optimize,
)
grad2 = np.einsum(
"mn, cim, dcn -> dc",
self.parameters["gcoeff"],
self.ao_val[:, e + 1 :, :],
e_grad[:, :, 0, :],
optimize=self.optimize,
)
grad = grad1 + grad2
lap3 = np.einsum("dc,dc->c", grad, grad)
return gpu.asnumpy(grad), gpu.asnumpy(lap1 + lap2 + lap3)
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),
}
def __call__(self, configs, wf):
""""""
nconf = configs.configs.shape[0]
if not self._warmed_up:
naux = nconf if self._naux is None else self._naux
self.warm_up(naux)
self._warmed_up = True
dtype = complex if self.iscomplex else float
results = {
"value": np.zeros((nconf, self.norb, self.norb), dtype=dtype),
"norm": np.zeros((nconf, self.norb)),
}
naux = self._extra_config.configs.shape[0]
auxassignments = np.random.randint(0, naux, size=(self._nsweeps, nconf))
accept, extra_configs, borb_aux = sample_onebody(
self._extra_config,
self.orbitals,
spin=0,
nsamples=self._nsweeps,
tstep=self._tstep,
)
self._extra_config = extra_configs[-1]
for conf, assign in zip(extra_configs, auxassignments):
conf.resample(assign)
borb_aux = cp.asarray(
[orb[assign, ...] for orb, assign in zip(borb_aux, auxassignments)]
)
borb_configs = self.evaluate_orbitals(configs.electron(self._electrons))
borb_configs = borb_configs.reshape(nconf, self.nelec, -1)
bauxsquared = cp.abs(borb_aux) ** 2
fsum = cp.sum(bauxsquared, axis=-1, keepdims=True)
norm = bauxsquared / fsum
baux_f = borb_aux / fsum
for sweep, aux in enumerate(auxassignments):
wfratio = wf.testvalue_many(
self._electrons, extra_configs[sweep].electron(0)
)
ratio = np.einsum(
"ie,ij,iek->ijk",
wfratio.conj(),
baux_f[sweep, :],
borb_configs.conj(),
optimize=True,
)
results["value"] += asnumpy(ratio)
results["norm"] += asnumpy(norm[sweep, :])
results["value"] /= self._nstep
results["norm"] = results["norm"] / self._nstep
return results
def compute_value(updets, dndets, det_coeffs):
"""
Given the up and down determinant values, safely compute the total log wave function.
"""
upref = gpu.cp.amax(updets[1]).real
dnref = gpu.cp.amax(dndets[1]).real
phases = updets[0] * dndets[0]
logvals = updets[1] - upref + dndets[1] - dnref
phases = updets[0] * dndets[0]
wf_val = gpu.cp.einsum("d,id->i", det_coeffs, phases * gpu.cp.exp(logvals))
wf_sign = np.nan_to_num(wf_val / gpu.cp.abs(wf_val))
wf_logval = np.nan_to_num(gpu.cp.log(gpu.cp.abs(wf_val)) + upref + dnref)
return gpu.asnumpy(wf_sign), gpu.asnumpy(wf_logval)
def opt_hdf(hdf_file, data, attr, configs, parameters):
import pyqmc.hdftools as hdftools
if hdf_file is not None:
with h5py.File(hdf_file, "a") as hdf:
if "configs" not in hdf.keys():
hdftools.setup_hdf(hdf, data, attr)
configs.initialize_hdf(hdf)
hdf.create_group("wf")
for k, it in parameters.items():
hdf.create_dataset("wf/" + k, data=gpu.asnumpy(it))
hdftools.append_hdf(hdf, data)
configs.to_hdf(hdf)
for k, it in parameters.items():
hdf["wf/" + k][...] = gpu.asnumpy(it.copy())
def sample_onebody(configs, orbitals, spin, nsamples=1, tstep=0.5):
r""" For a set of orbitals defined by orb_coeff, return samples from :math:`f(r) = \sum_i \phi_i(r)^2`. """
n = configs.configs.shape[0]
ao = orbitals.aos("GTOval_sph", configs)
borb = orbitals.mos(ao, spin=spin)
fsum = (cp.abs(borb) ** 2).sum(axis=1)
allaccept = np.zeros((nsamples, n))
allconfigs = []
allorbs = []
for s in range(nsamples):
shift = np.sqrt(tstep) * np.random.randn(*configs.configs.shape)
newconfigs = configs.make_irreducible(0, configs.configs + shift)
ao = orbitals.aos("GTOval_sph", newconfigs)
borbnew = orbitals.mos(ao, spin=spin)
fsumnew = (cp.abs(borbnew) ** 2).sum(axis=1)
accept = asnumpy(fsumnew / fsum) > np.random.rand(n)
configs.move_all(newconfigs, accept)
borb[accept] = borbnew[accept]
fsum[accept] = fsumnew[accept]
allconfigs.append(configs.copy())
allaccept[s] = accept
allorbs.append(borb.copy())
return allaccept, allconfigs, allorbs
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)
def gradient(self, e, epos):
r"""We compute the gradient for electron e as
:math:`\nabla_e \ln \Psi_J = \sum_l c_l \left(\sum_{j > e} \nabla_e b_l(r_{ej}) + \sum_{i < e} \nabla_e b_l(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 generate_slater(
mol,
mf,
optimize_determinants=False,
optimize_orbitals=False,
optimize_zeros=True,
epsilon=1e-8,
**kwargs,
):
"""Construct a Slater determinant
:parameter boolean optimize_orbitals: make `to_opt` true for orbital parameters
:parameter array-like twist: The twist to extract from the mean-field object
:parameter boolean optimize_zeros: optimize coefficients that are zero in the mean-field object
:returns: slater, to_opt
"""
wf = slater.Slater(mol, mf, **kwargs)
to_opt = {}
if optimize_determinants:
to_opt["det_coeff"] = np.ones_like(wf.parameters["det_coeff"],
dtype=bool)
to_opt["det_coeff"][np.argmax(wf.parameters["det_coeff"])] = False
if optimize_orbitals:
for k in ["mo_coeff_alpha", "mo_coeff_beta"]:
to_opt[k] = np.ones(wf.parameters[k].shape, dtype=bool)
if not optimize_zeros:
to_opt[k][
np.abs(gpu.asnumpy(wf.parameters[k])) < epsilon] = False
return wf, to_opt
def __init__(self, parameters, to_opt=None):
parameters = {k: gpu.asnumpy(v) for k, v in parameters.items()}
if to_opt is None:
to_opt = {
k: np.ones(p.shape, dtype=bool)
for k, p in parameters.items()
}
self.to_opt = {k: o for k, o in to_opt.items() if np.any(o)}
self.frozen_parms = {
k: parameters[k][~opt]
for k, opt in self.to_opt.items()
}
self.shapes = {k: parameters[k].shape for k in self.to_opt}
self.slices = {k: np.prod(s) for k, s in self.shapes.items()}
self.dtypes = {k: parameters[k].dtype for k in self.to_opt}
self.complex = {k: d == complex for k, d in self.dtypes.items()}
self.nimag = {
k: to_opt[k].sum() if c else 0
for k, c in self.complex.items()
}
self.complex_inds = np.concatenate([
np.ones(to_opt[k].sum(), dtype=bool) * c
for k, c in self.complex.items()
])
self.nparams = np.sum([v.sum() for v in self.to_opt.values()])
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)
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))
def laplacian(self, e, epos):
""" Compute the laplacian Psi/ Psi. """
s = int(e >= self._nelec[0])
ao = self.orbitals.aos("GTOval_sph_deriv2", epos)
ao_val = ao[:, 0, :, :]
ao_lap = gpu.cp.sum(ao[:, [4, 7, 9], :, :], axis=1)
mos = gpu.cp.stack(
[self.orbitals.mos(x, s)[..., self._det_occup[s]] for x in [ao_val, ao_lap]]
)
ratios = self._testrowderiv(e, mos)
return gpu.asnumpy(ratios[1] / ratios[0])
def testvalue(self, e, epos, mask=None):
"""return the ratio between the current wave function and the wave function if
electron e's position is replaced by epos"""
s = int(e >= self._nelec[0])
ao = self.orbitals.aos("GTOval_sph", epos, mask)
mo = self.orbitals.mos(ao, s)
mo_vals = mo[..., self._det_occup[s]]
if len(epos.configs.shape) > 2:
mo_vals = mo_vals.reshape(-1, epos.configs.shape[1],
mo_vals.shape[1], mo_vals.shape[2])
return gpu.asnumpy(self._testrow(e, mo_vals, mask))
def gradient_laplacian(self, e, epos):
s = int(e >= self._nelec[0])
ao = self.orbitals.aos("GTOval_sph_deriv2", epos)
ao = gpu.cp.concatenate(
[ao[:, 0:4, ...], ao[:, [4, 7, 9], ...].sum(axis=1, keepdims=True)], axis=1
)
mo = self.orbitals.mos(ao, s)
mo_vals = mo[:, :, self._det_occup[s]]
ratios = self._testrowderiv(e, mo_vals)
ratios = gpu.asnumpy(ratios / ratios[:1])
return ratios[1:-1], ratios[-1]
def gradient(self, e, epos):
"""Compute the gradient of the log wave function
Note that this can be called even if the internals have not been updated for electron e,
if epos differs from the current position of electron e."""
s = int(e >= self._nelec[0])
aograd = self.orbitals.aos("GTOval_sph_deriv1", epos)
mograd = self.orbitals.mos(aograd, s)
mograd_vals = mograd[:, :, self._det_occup[s]]
ratios = self._testrowderiv(e, mograd_vals)
return gpu.asnumpy(ratios[1:] / ratios[0])
def value(self):
mask = np.tril(np.ones((self.nelec, self.nelec)), -1)
vals = np.einsum(
"mn,cim, cjn, ij-> c",
self.parameters["gcoeff"],
self.ao_val,
self.ao_val,
mask,
optimize=self.optimize,
)
signs = np.ones(len(vals))
return (signs, gpu.asnumpy(vals))
def testvalue_many(self, e, epos, mask=None):
"""return the ratio between the current wave function and the wave function if
electron e's position is replaced by epos for each electron"""
s = (e >= self._nelec[0]).astype(int)
ao = self.orbitals.aos("GTOval_sph", epos, mask)
ratios = gpu.cp.zeros((epos.configs.shape[0], e.shape[0]), dtype=self.dtype)
for spin in [0, 1]:
ind = s == spin
mo = self.orbitals.mos(ao, spin)
mo_vals = mo[..., self._det_occup[spin]]
ratios[:, ind] = self._testrow(e[ind], mo_vals, mask, spin=spin)
return gpu.asnumpy(ratios)
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 gradient(self, e, epos):
_, e_grad = self._get_val_grad_lap(epos, mode="grad")
grad1 = np.einsum(
"mn, dcm, cjn -> dc",
self.parameters["gcoeff"],
e_grad[:, :, 0, :],
self.ao_val[:, :e, :],
optimize=self.optimize,
)
grad2 = np.einsum(
"mn, cim, dcn -> dc",
self.parameters["gcoeff"],
self.ao_val[:, e + 1 :, :],
e_grad[:, :, 0, :],
optimize=self.optimize,
)
return gpu.asnumpy(grad1 + grad2)
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
def testvalue(self, e, epos, mask=None):
if mask is None:
mask = [True] * epos.configs.shape[0]
masked_ao_val = self.ao_val[mask]
curr_val = np.einsum(
"mn, cm, cjn -> c",
self.parameters["gcoeff"],
masked_ao_val[:, e, :],
masked_ao_val[:, :e, :],
optimize=self.optimize,
)
curr_val += np.einsum(
"mn, cim, cn -> c",
self.parameters["gcoeff"],
masked_ao_val[:, e + 1 :, :],
masked_ao_val[:, e, :],
optimize=self.optimize,
)
new_ao_val = self._get_val_grad_lap(epos, mode="val", mask=mask)
new_val = np.einsum(
"mn, c...m, cjn -> c...",
self.parameters["gcoeff"],
new_ao_val,
masked_ao_val[:, :e, :],
optimize=self.optimize,
)
new_val += np.einsum(
"mn, cim, c...n -> c...",
self.parameters["gcoeff"],
masked_ao_val[:, e + 1 :, :],
new_ao_val,
optimize=self.optimize,
)
return gpu.asnumpy(np.exp((new_val.T - curr_val).T))
def serialize_parameters(self, parameters):
"""Convert the dictionary to a linear list of gradients"""
params = np.concatenate([
gpu.asnumpy(parameters[k])[opt] for k, opt in self.to_opt.items()
])
return np.concatenate((params.real, params[self.complex_inds].imag))
def __call__(self, configs, wf):
"""Gathers quantities from equation (10) of DOI:10.1063/1.4793531.
assignments maps from the auxilliary walkers onto the main walkers.
It should be of length [nsweeps,nconf], and contain integers between 0 and naux.
"""
nconf = configs.configs.shape[0]
if not self._warmed_up:
naux = nconf if self._naux is None else self._naux
self.warm_up(naux)
self._warmed_up = True
aux = self.get_configurations(nconf)
# Evaluate orbital values for the primary samples
ao_configs = self.orbitals.aos("GTOval_sph", configs)
ao_configs = ao_configs.reshape(
(ao_configs.shape[0], nconf, -1, ao_configs.shape[-1]))
orb_configs = [
self.orbitals.mos(ao_configs[..., self._electrons[spin], :], spin)
for spin in [0, 1]
]
results = {
"value": np.zeros((nconf, self._ijkl.shape[1]), dtype=self.dtype),
"norm_a": np.zeros((nconf, orb_configs[0].shape[-1])),
"norm_b": np.zeros((nconf, orb_configs[1].shape[-1])),
}
orb_configs = gpu.cp.asarray(
[orb_configs[s][:, :, self._ijkl[2 * s]] for s in [0, 1]])
down_start = [np.min(self._electrons[s]) for s in [0, 1]]
for sweep in range(self._nsweeps):
fsum = [
gpu.cp.sum(gpu.cp.abs(aux["orbs"][spin][sweep])**2, axis=1)
for spin in [0, 1]
]
norm = [
gpu.cp.abs(aux["orbs"][spin][sweep])**2 /
fsum[spin][:, np.newaxis] for spin in [0, 1]
]
wfratio = []
electrons_a_ind = []
electrons_b_ind = []
for ea in self._electrons[0]:
# Don't move the same electron twice
electrons_b = self._electrons[1][self._electrons[1] != ea]
epos_a = aux["configs"][0][sweep].electron(0)
epos_b = aux["configs"][1][sweep].electron(0)
wfratio_a = wf.testvalue(ea, epos_a)
wf.updateinternals(ea, epos_a, configs)
wfratio_b = wf.testvalue_many(electrons_b, epos_b)
wf.updateinternals(ea, configs.electron(ea), configs)
wfratio.append(wfratio_a[:, np.newaxis] * wfratio_b)
electrons_a_ind.extend([ea - down_start[0]] * len(electrons_b))
electrons_b_ind.extend(electrons_b - down_start[1])
wfratio = np.concatenate(wfratio, axis=1)
"""
orbratio collects
phi_i(r1) phi_j(r1') phi_k(r2) phi_l(r2')/rho(r1') rho(r2')
"""
phi_j_r1p = aux["orbs"][0][sweep][..., self._ijkl[1]]
phi_l_r2p = aux["orbs"][1][sweep][..., self._ijkl[3]]
rho1rho2 = 1.0 / (fsum[0] * fsum[1])
# n is the walker number, i is the electron pair index, o is the orbital
orbratio = gpu.cp.einsum(
"nio,nio,no,no,n ->nio",
orb_configs[0][:, electrons_a_ind, :], # phi_i(r1)
orb_configs[1][:, electrons_b_ind, :], # phi_k(r2)
phi_j_r1p, # phi_j
phi_l_r2p, # phi_l
rho1rho2,
)
results["value"] += gpu.asnumpy(
gpu.cp.einsum("in,inj->ij", wfratio, orbratio))
results["norm_a"] += gpu.asnumpy(norm[0])
results["norm_b"] += gpu.asnumpy(norm[1])
results["value"] /= self._nsweeps
for e in ["a", "b"]:
results["norm_%s" % e] /= self._nsweeps
return results
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))
def pgradient(self):
"""Compute the parameter gradient of Psi.
Returns :math:`\partial_p \Psi/\Psi` as a dictionary of numpy arrays,
which correspond to the parameter dictionary.
The wave function is given by ci Di, with an implicit sum
We have two sets of parameters:
Determinant coefficients:
di psi/psi = Dui Ddi/psi
Orbital coefficients:
dj psi/psi = ci dj (Dui Ddi)/psi
Let's suppose that j corresponds to an up orbital coefficient. Then
dj (Dui Ddi) = (dj Dui)/Dui Dui Ddi/psi = (dj Dui)/Dui di psi/psi
where di psi/psi is the derivative defined above.
"""
d = {}
# Det coeff
curr_val = self.value()
nonzero = curr_val[0] != 0.0
#dets[spin][ (phase,log), configuration, determinant]
dets = (self._dets[0][:, :, self._det_map[0]],
self._dets[1][:, :, self._det_map[1]])
d["det_coeff"] = np.zeros(dets[0].shape[1:], dtype=dets[0].dtype)
d["det_coeff"][nonzero, :] = (
dets[0][0, nonzero, :] * dets[1][0, nonzero, :] *
gpu.cp.exp(dets[0][1, nonzero, :] + dets[1][1, nonzero, :] -
gpu.cp.array(curr_val[1][nonzero, np.newaxis])) /
gpu.cp.array(curr_val[0][nonzero, np.newaxis]))
for s, parm in zip([0, 1], ["mo_coeff_alpha", "mo_coeff_beta"]):
ao = self._aovals[:, :, s * self._nelec[0]:self._nelec[s] +
s * self._nelec[0], :]
split, aos = self.orbitals.pgradient(ao, s)
mos = gpu.cp.split(gpu.cp.arange(split[-1]),
gpu.asnumpy(split).astype(int))
# Compute dj Diu/Diu
nao = aos[0].shape[-1]
nconf = aos[0].shape[0]
nmo = int(split[-1])
deriv = gpu.cp.zeros((len(self._det_occup[s]), nconf, nao, nmo),
dtype=curr_val[0].dtype)
for det, occ in enumerate(self._det_occup[s]):
for ao, mo in zip(aos, mos):
for i in mo:
if i in occ:
col = occ.index(i)
deriv[det, :, :,
i] = self._testcol(det, col, s, ao)
# now we reduce over determinants
d[parm] = gpu.cp.zeros(deriv.shape[1:], dtype=curr_val[0].dtype)
for di, coeff in enumerate(self.parameters["det_coeff"]):
whichdet = self._det_map[s][di]
d[parm] += (deriv[whichdet] * coeff *
d["det_coeff"][:, di, np.newaxis, np.newaxis])
for k, v in d.items():
d[k] = gpu.asnumpy(v)
for k in list(d.keys()):
if np.prod(d[k].shape) == 0:
del d[k]
return d
