""" Evaluate the OBDM for a wave function object. """
import pyqmc.orbitals
import numpy as np
import pyqmc.mc as mc
from pyqmc.gpu import cp, asnumpy
import pyqmc.supercell as supercell
[docs]class OBDMAccumulator:
r"""Return the obdm as an array with indices rho[spin][i][j] = <c_{spin,i}c^+_{spin,j}>
.. math:: \rho^\sigma_{ij} = \langle c_{\sigma, i} c^\dagger_{\sigma, j} \rangle
We are measuring the amplitude of moving one electron (e.g. the first one) from orbital :math:`\phi_i` to orbital :math:`\phi_j` (Eq (7) of DOI:10.1063/1.4793531)
.. math:: \rho_{i,k} = \int dR dr' \Psi^*(R') \phi_k(r') \phi_i^*(r) \Psi(R)
(The complex conjugate is on the wrong orbital in Eq (7) in the paper.) Sampling :math:`R` from :math:`|\Psi(R)^2|` and :math:`r'` from :math:`f(r) = \sum_i |\phi(r)|^2`
.. math:: \rho_{i,k} = \int dR dr' \frac{\Psi^*(R')}{\Psi^*(R)} \left[\Psi^*(R) \Psi(R)\right] \frac{\phi_k(r') \phi_i^*(r)}{f(r)} \left[f(r)\right]
The distributions (in square brackets) are accounted for by the Monte Carlo integration
.. math:: \rho_{i,k} = \left\langle \frac{\Psi^*(R')}{\Psi^*(R)} \frac{\phi_k(r') \phi_i^*(r)}{f(r)} \right\rangle
Eq (9) in the paper is the complex conjugate of this
.. math:: \rho_{i,k}^* = \left\langle \frac{\Psi(R')}{\Psi(R)} \frac{\phi_k^*(r') \phi_i(r)}{f(r)} \right\rangle
Args:
mol (Mole): PySCF Mole object.
configs (array): electron positions.
wf (pyqmc wave function object): wave function to evaluate on.
orb_coeff (array): coefficients with size (nbasis,norb) relating mol basis to basis
of 1-RDM desired.
tstep (float): width of the Gaussian to update a walker position for the
extra coordinate.
spin: 0 or 1 for up or down. Defaults to all electrons.
"""
def __init__(
self,
mol,
orb_coeff,
nsweeps=5,
tstep=0.50,
warmup=100,
naux=None,
spin=None,
electrons=None,
kpts=None,
):
if spin is not None:
if spin == 0:
self._electrons = np.arange(0, mol.nelec[0])
elif spin == 1:
self._electrons = np.arange(mol.nelec[0], np.sum(mol.nelec))
else:
raise ValueError("Spin not equal to 0 or 1")
elif electrons is not None:
self._electrons = electrons
else:
self._electrons = np.arange(0, np.sum(mol.nelec))
if kpts is None:
self.orbitals = pyqmc.orbitals.MoleculeOrbitalEvaluator(
mol, [orb_coeff, orb_coeff]
)
if hasattr(mol, "a"):
raise ValueError("kpts is required if the system is periodic")
else:
if not hasattr(mol, "original_cell"):
mol = supercell.get_supercell(mol, np.eye(3))
self.orbitals = pyqmc.orbitals.PBCOrbitalEvaluatorKpoints(
mol, [orb_coeff, orb_coeff], kpts
)
self.iscomplex = self.orbitals.iscomplex
self._tstep = tstep
self.nelec = len(self._electrons)
self._nsweeps = nsweeps
self._nstep = nsweeps * self.nelec
self._warmup = warmup
self._naux = naux
self._warmed_up = False
self._mol = mol
self.norb = self.orbitals.parameters["mo_coeff_alpha"].shape[-1]
def warm_up(self, naux):
self._extra_config = mc.initial_guess(self._mol, int(naux / self.nelec) + 1)
self._extra_config.reshape((-1, 1, 3))
self._extra_config.resample(range(naux))
_, extra_configs, _ = sample_onebody(
self._extra_config,
self.orbitals,
spin=0,
nsamples=self._warmup,
tstep=self._tstep,
)
self._extra_config = extra_configs[-1]
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 avg(self, configs, wf):
return {k: np.mean(it, axis=0) for k, it in self(configs, wf).items()}
def evaluate_orbitals(self, configs):
ao = self.orbitals.aos("GTOval_sph", configs)
return self.orbitals.mos(ao, spin=0)
def keys(self):
return set(
["value", "norm"],
)
def shapes(self):
norb = self.norb
return {
"value": (norb, norb),
"norm": (norb,),
}
[docs]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 normalize_obdm(obdm, norm):
return obdm / (norm[np.newaxis, :] * norm[:, np.newaxis]) ** 0.5