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