def test_pbc_wfs():
"""
Ensure that the wave function objects are consistent in several situations.
"""
from pyscf.pbc import lib, gto, scf
from pyqmc.supercell import get_supercell
from pyqmc.slater import Slater
from pyqmc.multiplywf import MultiplyWF
from pyqmc import default_jastrow
import pyqmc
mol = gto.M(
atom="H 0. 0. 0.; H 1. 1. 1.",
basis="sto-3g",
unit="bohr",
a=(np.ones((3, 3)) - np.eye(3)) * 4,
)
mf = scf.KRKS(mol, mol.make_kpts((2, 2, 2))).run()
# mf_rohf = scf.KROKS(mol).run()
# mf_uhf = scf.KUKS(mol).run()
epsilon = 1e-5
nconf = 10
supercell = get_supercell(mol, S=(np.ones((3, 3)) - 2 * np.eye(3)))
epos = pyqmc.initial_guess(supercell, nconf)
# For multislaterpbc
# kinds = 0, 3, 5, 6 # G, X, Y, Z
# d1 = {kind: [0] for kind in kinds}
# d2 = d1.copy()
# d2.update({0: [], 3: [0, 1]})
# detwt = [2 ** 0.5, 2 ** 0.5]
# occup = [[d1, d2], [d1]]
# map_dets = [[0, 1], [0, 0]]
for wf in [
MultiplyWF(Slater(supercell, mf), default_jastrow(supercell)[0]),
Slater(supercell, mf),
]:
for k in wf.parameters:
if "mo_coeff" not in k and k != "det_coeff":
wf.parameters[k] = np.random.rand(*wf.parameters[k].shape)
_, epos = pyqmc.vmc(wf, epos, nblocks=1, nsteps=2, tstep=1) # move off node
for fname, func in zip(
["gradient", "laplacian", "pgradient"],
[
testwf.test_wf_gradient,
testwf.test_wf_laplacian,
testwf.test_wf_pgradient,
],
):
err = []
for delta in [1e-4, 1e-5, 1e-6, 1e-7, 1e-8]:
err.append(func(wf, epos, delta))
print(type(wf), fname, min(err))
assert min(err) < epsilon
for k, item in testwf.test_updateinternals(wf, epos).items():
print(k, item)
assert item < epsilon
def from_mean_field(self, cell, mf, twist=None):
"""
mf is expected to be a KUHF, KRHF, or equivalent DFT objects.
Selects occupied orbitals from a given twist
If cell is a supercell, will automatically choose the folded k-points that correspond to that twist.
"""
cell = (cell if hasattr(cell, "original_cell") else get_supercell(
cell, np.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]])))
if twist is None:
twist = np.zeros(3)
else:
twist = np.dot(np.linalg.inv(cell.a), np.mod(twist,
1.0)) * 2 * np.pi
kinds = list(
set(get_k_indices(cell, mf,
get_supercell_kpts(cell) + twist)))
if len(kinds) != cell.scale:
raise ValueError(
"Did not find the right number of k-points for this supercell")
detcoeff = np.array([1.0])
det_map = np.array([[0], [0]])
if len(mf.mo_coeff[0][0].shape) == 2:
occup_k = [[[list(np.argwhere(mf.mo_occ[spin][k] > 0.5)[:, 0])]
for k in kinds] for spin in [0, 1]]
elif len(mf.mo_coeff[0][0].shape) == 1:
occup_k = [[[list(np.argwhere(mf.mo_occ[k] > 1.5 - spin)[:, 0])]
for k in kinds] for spin in [0, 1]]
occup = [[], []]
for spin in [0, 1]:
count = 0
for occ_k in occup_k[spin]:
occup[spin] += [o + count for o in occ_k[0]]
count += len(occ_k[0])
kpts = mf.kpts[kinds]
if len(mf.mo_coeff[0][0].shape) == 2:
mo_coeff = [[
mf.mo_coeff[spin][k][:, occup_k[spin][kinds.index(k)][0]]
for k in kinds
] for spin in [0, 1]]
elif len(mf.mo_coeff[0][0].shape) == 1:
mo_coeff = [[
mf.mo_coeff[k][:, occup_k[spin][kinds.index(k)][0]]
for k in kinds
] for spin in [0, 1]]
else:
raise ValueError("Did not expect an scf object of type", type(mf))
for s in [0, 1]:
occup[s] = [occup[s]]
return (
detcoeff,
occup,
det_map,
PBCOrbitalEvaluatorKpoints(cell, mo_coeff, kpts),
)
def get_ewald_energy(cell, S, configs):
supercell = get_supercell(cell, S)
ewald = pyqmc.ewald.Ewald(supercell)
configs = PeriodicConfigs(configs, supercell.lattice_vectors())
ee, ei, ii = ewald.energy(configs)
etot = ee + ei + ii
print(dict(ee=ee, ei=ei, ii=ii))
print("total energy", etot)
return etot
def __init__(
self,
mol,
orb_coeff,
spin,
nsweeps=4,
tstep=0.50,
warmup=200,
naux=None,
ijkl=None,
kpts=None,
):
self._mol = mol
self._tstep = tstep
self._nsweeps = nsweeps
self._spin = spin
self._naux = naux
self._warmup = warmup
if kpts is None:
self.orbitals = pyqmc.orbitals.MoleculeOrbitalEvaluator(mol, orb_coeff)
norb_up = orb_coeff[0].shape[1]
norb_down = orb_coeff[1].shape[1]
if hasattr(mol, "a"):
warnings.warn(
"Using molecular orbital evaluator for a periodic system. This is likely wrong unless you know what you're doing. Make sure to pass kpts into TBDM if you want to use the periodic orbital evaluator."
)
else:
if not hasattr(mol, "original_cell"):
mol = supercell.get_supercell(mol, np.eye(3))
self.orbitals = pyqmc.orbitals.PBCOrbitalEvaluatorKpoints(
mol, orb_coeff, kpts
)
norb_up = np.sum([o.shape[1] for o in orb_coeff[0]])
norb_down = np.sum([o.shape[1] for o in orb_coeff[1]])
self.dtype = complex if self.orbitals.iscomplex else float
self._spin_sector = spin
self._electrons = [
np.arange(spin[s] * mol.nelec[0], mol.nelec[0] + spin[s] * mol.nelec[1])
for s in [0, 1]
]
# Default to full 2rdm if ijkl not specified
if ijkl is None:
ijkl = [
[i, j, k, l]
for i in range(norb_up)
for j in range(norb_up)
for k in range(norb_down)
for l in range(norb_down)
]
self._ijkl = np.array(ijkl).T
self._warmed_up = False
def from_mean_field(self,
cell,
mf,
twist=None,
determinants=None,
tol=None):
"""
mf is expected to be a KUHF, KRHF, or equivalent DFT objects.
Selects occupied orbitals from a given twist
If cell is a supercell, will automatically choose the folded k-points that correspond to that twist.
"""
cell = (cell
if hasattr(cell, "original_cell") else supercell.get_supercell(
cell, np.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]])))
if twist is None:
twist = np.zeros(3)
else:
twist = np.dot(np.linalg.inv(cell.a), np.mod(twist,
1.0)) * 2 * np.pi
kinds = list(
set(
get_k_indices(cell, mf,
supercell.get_supercell_kpts(cell) + twist)))
if len(kinds) != cell.scale:
raise ValueError(
f"Found {len(kinds)} k-points but should have found {cell.scale}."
)
kpts = mf.kpts[kinds]
if determinants is None:
determinants = [
(1.0,
pyqmc.determinant_tools.create_pbc_determinant(cell, mf, []))
]
mo_coeff, determinants_flat = select_orbitals_kpoints(
determinants, mf, kinds)
detcoeff, occup, det_map = pyqmc.determinant_tools.create_packed_objects(
determinants_flat, format="list", tol=tol)
# Check
for s, (occ_s, nelec_s) in enumerate(zip(occup, cell.nelec)):
for determinant in occ_s:
if len(determinant) != nelec_s:
raise RuntimeError(
f"The number of electrons of spin {s} should be {nelec_s}, but found {len(determinant)} orbital[s]. You may have used a large smearing value.. Please pass your own determinants list. "
)
return (
detcoeff,
occup,
det_map,
PBCOrbitalEvaluatorKpoints(cell, mo_coeff, kpts),
)
def __init__(
self,
mol,
orb_coeff,
nsweeps=5,
tstep=0.50,
warmup=100,
naux=500,
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))
self.iscomplex = bool(sum(map(np.iscomplexobj, orb_coeff)))
if kpts is None:
self.orbitals = MoleculeOrbitalEvaluator(mol,
[orb_coeff, orb_coeff])
else:
if not hasattr(mol, "original_cell"):
mol = supercell.get_supercell(mol, np.eye(3))
self.orbitals = PBCOrbitalEvaluatorKpoints(mol,
[orb_coeff, orb_coeff],
kpts)
self._tstep = tstep
self.nelec = len(self._electrons)
self._nsweeps = nsweeps
self._nstep = nsweeps * self.nelec
self.norb = self.orbitals.parameters["mo_coeff_alpha"].shape[-1]
self._extra_config = initial_guess(mol, int(naux / self.nelec) + 1)
self._extra_config.reshape((-1, 1, 3))
accept, extra_configs, _ = sample_onebody(
self._extra_config,
self.orbitals,
spin=0,
nsamples=warmup,
tstep=self._tstep,
)
self._extra_config = extra_configs[-1]
def __init__(
self,
mol,
orb_coeff,
spin,
nsweeps=4,
tstep=0.50,
warmup=200,
naux=500,
ijkl=None,
kpts=None,
):
assert (
len(orb_coeff.shape) == 3
), "orb_coeff should be a list of orbital coefficients with size (2,num_mobasis,num_orb)."
self._mol = mol
self._tstep = tstep
self._nsweeps = nsweeps
self._spin = spin
if kpts is None:
self.orbitals = MoleculeOrbitalEvaluator(mol, orb_coeff)
else:
if not hasattr(mol, "original_cell"):
mol = supercell.get_supercell(mol, np.eye(3))
self.orbitals = PBCOrbitalEvaluatorKpoints(mol, orb_coeff, kpts)
self._spin_sector = spin
self._electrons = [
np.arange(spin[s] * mol.nelec[0],
mol.nelec[0] + spin[s] * mol.nelec[1]) for s in [0, 1]
]
# Initialization and warmup of configurations
nwalkers = int(naux / sum(self._mol.nelec))
self._aux_configs = []
for spin in [0, 1]:
self._aux_configs.append(initial_guess(mol, nwalkers))
self._aux_configs[spin].reshape((-1, 1, 3))
_, self._aux_configs[spin], _ = sample_onebody(
self._aux_configs[spin], self.orbitals, 0, nsamples=warmup)
self._aux_configs[spin] = self._aux_configs[spin][-1]
# Default to full 2rdm if ijkl not specified
if ijkl is None:
norb_up = orb_coeff[0].shape[1]
norb_down = orb_coeff[1].shape[1]
ijkl = [[i, j, k, l] for i in range(norb_up)
for j in range(norb_up) for k in range(norb_down)
for l in range(norb_down)]
self._ijkl = np.array(ijkl).T
def test_pbc(li_cubic_ccecp):
from pyqmc import supercell
import scipy
mol, mf = li_cubic_ccecp
# S = np.ones((3, 3)) - np.eye(3)
S = np.identity(3)
mol = supercell.get_supercell(mol, S)
kpts = supercell.get_supercell_kpts(mol)[:2]
kdiffs = mf.kpts[np.newaxis] - kpts[:, np.newaxis]
kinds = np.nonzero(np.linalg.norm(kdiffs, axis=-1) < 1e-12)[1]
# Lowdin orthogonalized AO basis.
# lowdin = lo.orth_ao(mol, "lowdin")
loiao = lo.iao.iao(mol.original_cell, mf.mo_coeff, kpts=kpts)
occs = [mf.mo_occ[k] for k in kinds]
coefs = [mf.mo_coeff[k] for k in kinds]
ovlp = mf.get_ovlp()[kinds]
lowdin = [lo.vec_lowdin(l, o) for l, o in zip(loiao, ovlp)]
lreps = [np.linalg.multi_dot([l.T, o, c]) for l, o, c in zip(lowdin, ovlp, coefs)]
# make AO to localized orbital coefficients.
mfobdm = [np.einsum("ij,j,kj->ik", l.conj(), o, l) for l, o in zip(lreps, occs)]
### Test OBDM calculation.
nconf = 500
nsteps = 100
warmup = 6
wf = Slater(mol, mf)
configs = initial_guess(mol, nconf)
obdm_dict = dict(mol=mol, orb_coeff=lowdin, kpts=kpts, nsweeps=4, warmup=10)
obdm = OBDMAccumulator(**obdm_dict)
df, coords = vmc(
wf,
configs,
nsteps=nsteps,
accumulators={"obdm": obdm}, # , "obdm_up": obdm_up, "obdm_down": obdm_down},
verbose=True,
)
obdm_est = {}
for k in ["obdm"]: # , "obdm_up", "obdm_down"]:
avg_norm = np.mean(df[k + "norm"][warmup:], axis=0)
avg_obdm = np.mean(df[k + "value"][warmup:], axis=0)
obdm_est[k] = normalize_obdm(avg_obdm, avg_norm)
mfobdm = scipy.linalg.block_diag(*mfobdm)
mae = np.mean(np.abs(obdm_est["obdm"] - mfobdm))
assert mae < 0.05, f"mae {mae}"
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 test_RKS(kind=0, nk=(1, 1, 1)):
L = 2
mol = gto.M(
atom="""He {0} {0} {0}""".format(0.0),
basis="sto-3g",
a=np.eye(3) * L,
unit="bohr",
)
kpts = mol.make_kpts(nk)
mf = scf.KRKS(mol, kpts)
mf.xc = "pbe"
# mf = mf.density_fit()
mf = mf.run()
supercell = get_supercell(mol, np.diag(nk))
runtest(supercell, mf, kind=kind)
def noncubic(kind=0, nk=(1, 1, 1)):
L = 3
mol = gto.M(
atom="""H {0} {0} {0}
H {1} {1} {1}""".format(0.0, L / 4),
basis="sto-3g",
a=(np.ones((3, 3)) - np.eye(3)) * L / 2,
spin=0,
unit="bohr",
)
kpts = mol.make_kpts(nk)
mf = scf.KUKS(mol, kpts)
mf.xc = "pbe"
# mf = mf.density_fit()
mf = mf.run()
supercell = get_supercell(mol, np.diag(nk))
runtest(supercell, mf, kind=kind)
def from_mean_field(self, cell, mf, twist=None, determinants=None):
"""
mf is expected to be a KUHF, KRHF, or equivalent DFT objects.
Selects occupied orbitals from a given twist
If cell is a supercell, will automatically choose the folded k-points that correspond to that twist.
"""
cell = (
cell
if hasattr(cell, "original_cell")
else supercell.get_supercell(
cell, np.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
)
)
if twist is None:
twist = np.zeros(3)
else:
twist = np.dot(np.linalg.inv(cell.a), np.mod(twist, 1.0)) * 2 * np.pi
kinds = list(
set(get_k_indices(cell, mf, supercell.get_supercell_kpts(cell) + twist))
)
if len(kinds) != cell.scale:
raise ValueError(
f"Found {len(kinds)} k-points but should have found {cell.scale}."
)
kpts = mf.kpts[kinds]
if determinants is None:
determinants = [
(1.0, pyqmc.determinant_tools.create_pbc_determinant(cell, mf, []))
]
mo_coeff, determinants_flat = select_orbitals_kpoints(determinants, mf, kinds)
detcoeff, occup, det_map = pyqmc.determinant_tools.create_packed_objects(
determinants_flat, format="list"
)
return (
detcoeff,
occup,
det_map,
PBCOrbitalEvaluatorKpoints(cell, mo_coeff, kpts),
)
def initialize_qmc_objects(
dft_checkfile,
nconfig=1000,
load_parameters=None,
ci_checkfile=None,
S=None,
jastrow_kws=None,
slater_kws=None,
accumulators=None,
opt_wf=False,
target_root=0,
nodal_cutoff=1e-3,
):
if ci_checkfile is None:
mol, mf = pyscftools.recover_pyscf(dft_checkfile)
mc = None
else:
mol, mf, mc = pyscftools.recover_pyscf(dft_checkfile,
ci_checkfile=ci_checkfile)
mc.ci = mc.ci[target_root]
if S is not None:
mol = supercell.get_supercell(mol, np.asarray(S))
wf, to_opt = wftools.generate_wf(mol,
mf,
mc=mc,
jastrow_kws=jastrow_kws,
slater_kws=slater_kws)
if load_parameters is not None:
wftools.read_wf(wf, load_parameters)
configs = pyqmc.mc.initial_guess(mol, nconfig)
if opt_wf:
acc = pyqmc.accumulators.gradient_generator(mol,
wf,
to_opt,
nodal_cutoff=nodal_cutoff)
else:
if accumulators == None:
accumulators = {}
acc = generate_accumulators(mol, mf, **accumulators)
return wf, configs, acc
def multislater(kind=0, nk=(1, 1, 1)):
L = 3
mol = gto.Cell(
atom="""H {0} {0} {0}
H {1} {1} {1}""".format(0.0, L / 2),
basis="cc-pvtz",
spin=0,
unit="bohr",
)
mol.exp_to_discard = 0.1
mol.build(a=np.eye(3) * L)
kpts = mol.make_kpts(nk)
mf = scf.UKS(mol, (0, 0, 0))
mf.xc = "pbe"
mf = multigrid(mf)
mf = remove_linear_dep_(mf)
mf.chkfile = "h_bcc.chkfile"
mf = mf.run()
supercell = get_supercell(mol, np.diag(nk))
runtest(supercell, mf, kind=kind, do_mc=True)
def test_info_functions_pbc():
from pyscf.pbc import gto, scf
from pyqmc.supercell import get_supercell
mol = gto.Cell(atom="He 0.00 0.00 0.00", basis="ccpvdz", unit="B")
mol.a = 5.61 * np.eye(3)
mol.build()
mf = scf.KRHF(mol, kpts=mol.make_kpts([2, 2, 2])).density_fit()
ehf = mf.kernel()
supercell = get_supercell(mol, 2 * np.eye(3))
kinds = [0, 1]
dm_orbs = [mf.mo_coeff[i][:, :2] for i in kinds]
wf, to_opt = pyqmc.default_sj(mol, mf)
accumulators = {
"pgrad": pyqmc.gradient_generator(mol, wf, to_opt, ewald_gmax=10),
"obdm": OBDMAccumulator(mol, dm_orbs, kpts=mf.kpts[kinds]),
"Sq": pyqmc.accumulators.SqAccumulator(mol.lattice_vectors()),
}
info_functions(mol, wf, accumulators)
def cubic_with_ecp(kind=0, nk=(1, 1, 1)):
from pyscf.pbc.dft.multigrid import multigrid
start = time.time()
L = 6.63
mol = gto.Cell(
atom="""Li {0} {0} {0}
Li {1} {1} {1}""".format(0.0, L / 2),
basis="bfd-vdz",
ecp="bfd",
spin=0,
unit="bohr",
)
mol.exp_to_discard = 0.1
mol.build(a=np.eye(3) * L)
kpts = mol.make_kpts(nk)
mf = scf.KUKS(mol, kpts)
mf.xc = "pbe"
# mf = mf.density_fit()
mf = multigrid(mf)
mf = mf.run()
supercell = get_supercell(mol, np.diag(nk))
runtest(supercell, mf, kind=kind)
kmf = scf.KRHF(cell, exxdiv=None).density_fit()
kmf.kpts = cell.make_kpts([nk, nk, nk])
ehf = kmf.kernel()
print("EHF", ehf)
return cell, kmf
if __name__ == "__main__":
# Run SCF
cell, kmf = run_scf(nk=2)
# Set up wf and configs
nconfig = 100
S = np.eye(3) * 2 # 2x2x2 supercell
supercell = get_supercell(cell, S)
wf, to_opt = pyqmc.default_sj(supercell, kmf)
configs = pyqmc.initial_guess(supercell, nconfig)
# Initialize energy accumulator (and Ewald)
pgrad = pyqmc.gradient_generator(supercell, wf, to_opt=to_opt)
# Optimize jastrow
wf, lm_df = pyqmc.line_minimization(wf,
configs,
pgrad,
hdf_file="pbc_he_linemin.hdf",
verbose=True)
# Run VMC
df, configs = pyqmc.vmc(
def __init__(
self,
mol,
orb_coeff,
nsweeps=5,
tstep=0.50,
warmup=100,
naux=500,
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))
self.iscomplex = bool(sum(map(np.iscomplexobj, orb_coeff)))
if hasattr(mol, "lattice_vectors"):
assert kpts is not None
assert len(orb_coeff) == len(kpts)
self.iscomplex = self.iscomplex or np.linalg.norm(kpts) > 1e-12
self._kpts = kpts
self.evaluate_orbitals = self._evaluate_orbitals_pbc
if self.iscomplex:
self.get_wrapphase = slater.get_wrapphase_complex
else:
self.get_wrapphase = slater.get_wrapphase_real
for attribute in ["original_cell", "S"]:
if not hasattr(mol, attribute):
from pyqmc.supercell import get_supercell
mol = get_supercell(mol, np.eye(3))
self.supercell = mol
self._mol = mol.original_cell
else:
self._mol = mol
self.evaluate_orbitals = self._evaluate_orbitals
assert (len(orb_coeff.shape) == 2
), "orb_coeff should be a list of orbital coefficients."
self._orb_coeff = orb_coeff
if hasattr(self._orb_coeff, "shape"):
self.norb = self._orb_coeff.shape[1]
else:
self.norb = sum(c.shape[1] for c in self._orb_coeff)
self._tstep = tstep
self.nelec = len(self._electrons)
self._nsweeps = nsweeps
self._nstep = nsweeps * self.nelec
self._extra_config = initial_guess(mol, int(naux / self.nelec) + 1)
self._extra_config.reshape((-1, 3))
accept, extra_configs = self.sample_onebody(self._extra_config, warmup)
self._extra_config = extra_configs[-1]
def test_pbc():
from pyscf.pbc import gto, scf
from pyqmc import supercell
import scipy
lvecs = (np.ones((3, 3)) - 2 * np.eye(3)) * 2.0
mol = gto.M(
atom="H 0. 0. -{0}; H 0. 0. {0}".format(0.7),
basis="sto-3g",
unit="bohr",
verbose=0,
a=lvecs,
)
mf = scf.KRHF(mol, kpts=mol.make_kpts((2, 2, 2)))
mf = mf.run()
S = np.ones((3, 3)) - np.eye(3)
mol = supercell.get_supercell(mol, S)
kpts = supercell.get_supercell_kpts(mol)[:2]
kdiffs = mf.kpts[np.newaxis] - kpts[:, np.newaxis]
kinds = np.nonzero(np.linalg.norm(kdiffs, axis=-1) < 1e-12)[1]
# Lowdin orthogonalized AO basis.
# lowdin = lo.orth_ao(mol, "lowdin")
loiao = lo.iao.iao(mol.original_cell, mf.mo_coeff, kpts=kpts)
occs = [mf.mo_occ[k] for k in kinds]
coefs = [mf.mo_coeff[k] for k in kinds]
ovlp = mf.get_ovlp()[kinds]
lowdin = [lo.vec_lowdin(l, o) for l, o in zip(loiao, ovlp)]
lreps = [
np.linalg.multi_dot([l.T, o, c])
for l, o, c in zip(lowdin, ovlp, coefs)
]
# make AO to localized orbital coefficients.
mfobdm = [
np.einsum("ij,j,kj->ik", l.conj(), o, l) for l, o in zip(lreps, occs)
]
### Test OBDM calculation.
nconf = 800
nsteps = 50
warmup = 6
wf = PySCFSlater(mol, mf)
configs = initial_guess(mol, nconf)
obdm_dict = dict(mol=mol,
orb_coeff=lowdin,
kpts=kpts,
nsweeps=4,
warmup=10)
obdm = OBDMAccumulator(**obdm_dict)
obdm_up = OBDMAccumulator(**obdm_dict, spin=0)
obdm_down = OBDMAccumulator(**obdm_dict, spin=1)
df, coords = vmc(
wf,
configs,
nsteps=nsteps,
accumulators={
"obdm": obdm,
"obdm_up": obdm_up,
"obdm_down": obdm_down
},
verbose=True,
)
obdm_est = {}
for k in ["obdm", "obdm_up", "obdm_down"]:
avg_norm = np.mean(df[k + "norm"][warmup:], axis=0)
avg_obdm = np.mean(df[k + "value"][warmup:], axis=0)
obdm_est[k] = normalize_obdm(avg_obdm, avg_norm)
print("Average OBDM(orb,orb)", obdm_est["obdm"].round(3))
mfobdm = scipy.linalg.block_diag(*mfobdm)
print("mf obdm", mfobdm.round(3))
max_abs_err = np.max(np.abs(obdm_est["obdm"] - mfobdm))
assert max_abs_err < 0.05, "max abs err {0}".format(max_abs_err)
print(obdm_est["obdm_up"].diagonal().round(3))
print(obdm_est["obdm_down"].diagonal().round(3))
mae = np.mean(np.abs(obdm_est["obdm_up"] + obdm_est["obdm_down"] - mfobdm))
maup = np.mean(np.abs(obdm_est["obdm_up"]))
madn = np.mean(np.abs(obdm_est["obdm_down"]))
mamf = np.mean(np.abs(mfobdm))
assert mae < 0.05, "mae {0}\n maup {1}\n madn {2}\n mamf {3}".format(
mae, maup, madn, mamf)
