def test_single_opt():
from pyqmc.accumulators import EnergyAccumulator
from pyscf import lib, gto, scf
import pandas as pd
from pyqmc.multiplywf import MultiplyWF
from pyqmc.jastrow import Jastrow2B
from pyqmc.func3d import GaussianFunction
from pyqmc.slater import PySCFSlaterRHF
from pyqmc.multiplywf import MultiplyWF
from pyqmc.jastrow import Jastrow2B
from pyqmc.mc import initial_guess,vmc
mol = gto.M(atom='Li 0. 0. 0.; Li 0. 0. 1.5', basis='bfd_vtz',ecp='bfd',unit='bohr',verbose=5)
mf = scf.RHF(mol).run()
nconf=1000
nsteps=10
coords = initial_guess(mol,nconf)
wf=MultiplyWF(PySCFSlaterRHF(mol,mf),Jastrow2B(mol,
basis=[GaussianFunction(1.0),GaussianFunction(2.0)]))
vmc(wf,coords,nsteps=nsteps)
opt_var,wf=optvariance(EnergyAccumulator(mol),wf,coords,['wf2coeff'])
print('Final variance:',opt_var)
def test_vmc():
"""
Test that a VMC calculation of a Slater determinant matches Hartree-Fock within error bars.
"""
nconf = 1000
mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5",
basis="cc-pvtz",
unit="bohr",
verbose=1)
mf_rhf = scf.RHF(mol).run()
mf_uhf = scf.UHF(mol).run()
nsteps = 300
warmup = 30
for wf, mf in [
(PySCFSlater(mol, mf_rhf), mf_rhf),
(PySCFSlater(mol, mf_uhf), mf_uhf),
]:
# Without blocks
coords = initial_guess(mol, nconf)
df, coords = vmc(
wf,
coords,
nsteps=nsteps,
accumulators={"energy": EnergyAccumulator(mol)},
verbose=True,
)
df = pd.DataFrame(df)
df = reblock(df["energytotal"][warmup:], 20)
en = df.mean()
err = df.sem()
assert en - mf.energy_tot(
) < 5 * err, "pyscf {0}, vmc {1}, err {2}".format(
mf.energy_tot(), en, err)
# With blocks
coords = initial_guess(mol, nconf)
df, coords = vmc(
wf,
coords,
nblocks=int(nsteps / 30),
nsteps_per_block=30,
accumulators={"energy": EnergyAccumulator(mol)},
)
df = pd.DataFrame(df)["energytotal"][int(warmup / 30):]
en = df.mean()
err = df.sem()
assert en - mf.energy_tot(
) < 5 * err, "pyscf {0}, vmc {1}, err {2}".format(
mf.energy_tot(), en, err)
def test_vmc():
"""
Test that a VMC calculation of a Slater determinant matches Hartree-Fock within error bars.
"""
nconf = 5000
mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5",
basis="cc-pvtz",
unit="bohr",
verbose=1)
mf_rhf = scf.RHF(mol).run()
mf_uhf = scf.UHF(mol).run()
nsteps = 100
warmup = 30
for wf, mf in [
(PySCFSlaterUHF(mol, mf_rhf), mf_rhf),
(PySCFSlaterUHF(mol, mf_uhf), mf_uhf),
]:
coords = initial_guess(mol, nconf)
df, coords = vmc(wf,
coords,
nsteps=nsteps,
accumulators={"energy": EnergyAccumulator(mol)})
df = pd.DataFrame(df)
en = np.mean(df["energytotal"][warmup:])
err = np.std(df["energytotal"][warmup:]) / np.sqrt(nsteps - warmup)
assert en - mf.energy_tot() < 10 * err
def test_accumulator():
""" Tests that the accumulator gets inserted into the data output correctly.
"""
import pandas as pd
from pyqmc.mc import vmc, initial_guess
from pyscf import gto, scf
from pyqmc.energy import energy
from pyqmc.slateruhf import PySCFSlaterUHF
from pyqmc.accumulators import EnergyAccumulator
mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5",
basis="cc-pvtz",
unit="bohr",
verbose=5)
mf = scf.RHF(mol).run()
nconf = 5000
wf = PySCFSlaterUHF(mol, mf)
coords = initial_guess(mol, nconf)
df, coords = vmc(wf,
coords,
nsteps=30,
accumulators={"energy": EnergyAccumulator(mol)})
df = pd.DataFrame(df)
eaccum = EnergyAccumulator(mol)
eaccum_energy = eaccum(coords, wf)
df = pd.DataFrame(df)
print(df["energytotal"][29] == np.average(eaccum_energy["total"]))
assert df["energytotal"][29] == np.average(eaccum_energy["total"])
def test():
""" Ensure that DMC obtains the exact result for a hydrogen atom """
from pyscf import lib, gto, scf
from pyqmc.slater import PySCFSlater
from pyqmc.jastrowspin import JastrowSpin
from pyqmc.dmc import limdrift, rundmc
from pyqmc.mc import vmc
from pyqmc.accumulators import EnergyAccumulator
from pyqmc.func3d import CutoffCuspFunction
from pyqmc.multiplywf import MultiplyWF
from pyqmc.coord import OpenConfigs
import pandas as pd
mol = gto.M(atom="H 0. 0. 0.", basis="sto-3g", unit="bohr", spin=1)
mf = scf.UHF(mol).run()
nconf = 1000
configs = OpenConfigs(np.random.randn(nconf, 1, 3))
wf1 = PySCFSlater(mol, mf)
wf = wf1
wf2 = JastrowSpin(mol, a_basis=[CutoffCuspFunction(5, 0.2)], b_basis=[])
wf2.parameters["acoeff"] = np.asarray([[[1.0, 0]]])
wf = MultiplyWF(wf1, wf2)
dfvmc, configs_ = vmc(
wf, configs, nsteps=50, accumulators={"energy": EnergyAccumulator(mol)}
)
dfvmc = pd.DataFrame(dfvmc)
print(
"vmc energy",
np.mean(dfvmc["energytotal"]),
np.std(dfvmc["energytotal"]) / np.sqrt(len(dfvmc)),
)
warmup = 200
branchtime = 5
dfdmc, configs_, weights_ = rundmc(
wf,
configs,
nsteps=4000 + warmup * branchtime,
branchtime=branchtime,
accumulators={"energy": EnergyAccumulator(mol)},
ekey=("energy", "total"),
tstep=0.01,
drift_limiter=limdrift,
verbose=True,
)
dfdmc = pd.DataFrame(dfdmc)
dfdmc.sort_values("step", inplace=True)
dfprod = dfdmc[dfdmc.step >= warmup]
rb_summary = reblock.reblock_summary(dfprod[["energytotal", "energyei"]], 20)
print(rb_summary)
energy, err = [rb_summary[v]["energytotal"] for v in ("mean", "standard error")]
assert (
np.abs(energy + 0.5) < 5 * err
), "energy not within {0} of -0.5: energy {1}".format(5 * err, np.mean(energy))
def test():
mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5",
basis="sto-3g",
unit="bohr",
verbose=0)
mf = scf.RHF(mol).run()
# Lowdin orthogonalized AO basis.
lowdin = lo.orth_ao(mol, "lowdin")
# MOs in the Lowdin basis.
mo = solve(lowdin, mf.mo_coeff)
# make AO to localized orbital coefficients.
mfobdm = mf.make_rdm1(mo, mf.mo_occ)
### Test OBDM calculation.
nconf = 500
nsteps = 400
obdm_steps = 4
warmup = 15
wf = PySCFSlaterUHF(mol, mf)
configs = initial_guess(mol, nconf)
energy = EnergyAccumulator(mol)
obdm = OBDMAccumulator(mol=mol, orb_coeff=lowdin, nsweeps=1)
obdm_up = OBDMAccumulator(mol=mol, orb_coeff=lowdin, nsweeps=1, spin=0)
obdm_down = OBDMAccumulator(mol=mol, orb_coeff=lowdin, nsweeps=1, spin=1)
df, coords = vmc(
wf,
configs,
nsteps=nsteps,
accumulators={
"energy": energy,
"obdm": obdm,
"obdm_up": obdm_up,
"obdm_down": obdm_down,
},
)
df = DataFrame(df)
obdm_est = {}
for k in ["obdm", "obdm_up", "obdm_down"]:
avg_norm = np.array(df.loc[warmup:,
k + "norm"].values.tolist()).mean(axis=0)
avg_obdm = np.array(df.loc[warmup:,
k + "value"].values.tolist()).mean(axis=0)
obdm_est[k] = normalize_obdm(avg_obdm, avg_norm)
print("Average OBDM(orb,orb)", obdm_est["obdm"].diagonal().round(3))
print("mf obdm", mfobdm.diagonal().round(3))
assert np.max(np.abs(obdm_est["obdm"] - mfobdm)) < 0.05
print(obdm_est["obdm_up"].diagonal().round(3))
print(obdm_est["obdm_down"].diagonal().round(3))
assert np.mean(
np.abs(obdm_est["obdm_up"] + obdm_est["obdm_down"] - mfobdm)) < 0.05
def test():
""" Ensure that DMC obtains the exact result for a hydrogen atom """
from pyscf import lib, gto, scf
from pyqmc.slateruhf import PySCFSlaterUHF
from pyqmc.jastrowspin import JastrowSpin
from pyqmc.dmc import limdrift, dmc
from pyqmc.mc import vmc
from pyqmc.accumulators import EnergyAccumulator
from pyqmc.func3d import ExpCuspFunction
from pyqmc.multiplywf import MultiplyWF
import pandas as pd
mol = gto.M(atom='H 0. 0. 0.', basis='sto-3g', unit='bohr', spin=1)
mf = scf.UHF(mol).run()
nconf = 1000
configs = np.random.randn(nconf, 1, 3)
wf1 = PySCFSlaterUHF(mol, mf)
wf = wf1
wf2 = JastrowSpin(mol, a_basis=[ExpCuspFunction(5, .2)], b_basis=[])
wf2.parameters['acoeff'] = np.asarray([[-1.0, 0]])
wf = MultiplyWF(wf1, wf2)
dfvmc, configs_ = vmc(wf,
configs,
nsteps=50,
accumulators={'energy': EnergyAccumulator(mol)})
dfvmc = pd.DataFrame(dfvmc)
print('vmc energy', np.mean(dfvmc['energytotal']),
np.std(dfvmc['energytotal']) / np.sqrt(len(dfvmc)))
dfdmc, configs_, weights_ = dmc(
wf,
configs,
nsteps=5000,
branchtime=5,
accumulators={'energy': EnergyAccumulator(mol)},
ekey=('energy', 'total'),
tstep=0.01,
drift_limiter=limdrift,
verbose=True)
dfdmc = pd.DataFrame(dfdmc)
dfdmc.sort_values('step', inplace=True)
warmup = 200
dfprod = dfdmc[dfdmc.step > warmup]
reblock = pyblock.reblock(dfprod[['energytotal', 'energyei']])
print(reblock[1])
dfoptimal = reblock[1][reblock[1][('energytotal', 'optimal block')] != '']
energy = dfoptimal[('energytotal', 'mean')].values[0]
err = dfoptimal[('energytotal', 'standard error')].values[0]
print("energy", energy, "+/-", err)
assert np.abs(
energy +
0.5) < 5 * err, "energy not within {0} of -0.5: energy {1}".format(
5 * err, np.mean(energy))
def test_updateinternals(wf, configs):
"""
:parameter wf: a wave function object to be tested
:parameter configs: electron positions
:type configs: (nconf, nelec, 3) array
:returns: max abs errors
:rtype: dictionary
"""
import pyqmc.mc as mc
nconf, ne, ndim = configs.configs.shape
delta = 1e-2
iscomplex = 1j if wf.iscomplex else 1
updatevstest = np.zeros((ne, nconf)) * iscomplex
recomputevstest = np.zeros((ne, nconf)) * iscomplex
recomputevsupdate = np.zeros((ne, nconf)) * iscomplex
wfcopy = copy.copy(wf)
val1 = wf.recompute(configs)
for e in range(ne):
# val1 = wf.recompute(configs)
epos = configs.make_irreducible(e, configs.configs[:, e, :] + delta)
ratio = wf.testvalue(e, epos)
wf.updateinternals(e, epos, configs)
update = wf.value()
configs.move(e, epos, [True] * nconf)
recompute = wfcopy.recompute(configs)
updatevstest[
e, :] = update[0] / val1[0] * np.exp(update[1] - val1[1]) - ratio
recomputevsupdate[e, :] = update[0] / val1[0] * np.exp(
update[1] -
val1[1]) - recompute[0] / val1[0] * np.exp(recompute[1] - val1[1])
recomputevstest[e, :] = (
recompute[0] / val1[0] * np.exp(recompute[1] - val1[1]) - ratio)
val1 = recompute
# Test mask and pgrad
_, configs = mc.vmc(wf, configs, nblocks=1, nsteps_per_block=1, tstep=2)
pgradupdate = wf.pgradient()
wf.recompute(configs)
pgrad = wf.pgradient()
pgdict = {
k: np.max(np.abs(pgu - pgrad[k]))
for k, pgu in pgradupdate.items() if np.prod(pgu.shape) > 0
}
return {
"updatevstest": np.max(np.abs(updatevstest)),
"recomputevstest": np.max(np.abs(recomputevstest)),
"recomputevsupdate": np.max(np.abs(recomputevsupdate)),
**pgdict,
}
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 test_ecp():
mol = gto.M(atom='C 0. 0. 0.', ecp='bfd', basis='bfd_vtz')
mf = scf.RHF(mol).run()
nconf=5000
wf=PySCFSlaterUHF(mol,mf)
coords = initial_guess(mol,nconf)
df,coords=vmc(wf,coords,nsteps=100,accumulators={'energy':EnergyAccumulator(mol)} )
df=pd.DataFrame(df)
warmup=30
print('mean field',mf.energy_tot(),'vmc estimation', np.mean(df['energytotal'][warmup:]),np.std(df['energytotal'][warmup:]))
assert abs(mf.energy_tot()-np.mean(df['energytotal'][warmup:])) <= np.std(df['energytotal'][warmup:])
def test():
mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5",
basis="sto-3g",
unit="bohr",
verbose=0)
mf = scf.RHF(mol).run()
# Lowdin orthogonalized AO basis.
lowdin = lo.orth_ao(mol, "lowdin")
# MOs in the Lowdin basis.
mo = solve(lowdin, mf.mo_coeff)
# make AO to localized orbital coefficients.
mfobdm = mf.make_rdm1(mo, mf.mo_occ)
### Test OBDM calculation.
nconf = 500
nsteps = 400
warmup = 15
wf = Slater(mol, mf)
configs = initial_guess(mol, nconf)
obdm_dict = dict(mol=mol, orb_coeff=lowdin, nsweeps=5, warmup=15)
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
},
)
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)
assert np.mean(
np.abs(obdm_est["obdm_up"] + obdm_est["obdm_down"] - mfobdm)) < 0.05
def genconfigs(n):
"""
Generate configurations and weights corresponding
to the highest determinant in MD expansion
"""
import os
os.system('mkdir -p vmc/')
mol, mf, mc, wf, to_opt, freeze = wavefunction(return_mf=True)
#Sample from the wave function which we're taking pderiv relative to
mf.mo_coeff = mc.mo_coeff
mf.mo_occ *= 0
mf.mo_occ[wf.wf1._det_occup[0][-1]] = 2
wfp = PySCFSlaterUHF(mol, mf)
#Lots of configurations
coords = pyqmc.initial_guess(mol, 100000)
eacc = EnergyAccumulator(mol)
transform = LinearTransform(wf.parameters, to_opt, freeze)
pgrad_bare = PGradTransform(eacc, transform, 0)
#Lots of steps
warmup = 10
for i in range(n + warmup + 1):
df, coords = vmc(wfp, coords, nsteps=1)
print(i)
if (i > warmup):
coords.configs.dump('vmc/coords' + str(i - warmup) + '.pickle')
val = wf.recompute(coords)
valp = wfp.value()
d = pgrad_bare(coords, wf)
data = {
'dpH': np.array(d['dpH'])[:, -1],
'dppsi': np.array(d['dppsi'])[:, -1],
'en': np.array(d['total']),
"wfval": val[1],
"wfpval": valp[1]
}
pd.DataFrame(data).to_json('vmc/evals' + str(i - warmup) + '.json')
return -1
def vmcparsl(wf, lastrun, nsteps, accumulators, stepoffset=0):
import os
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
from pyqmc.mc import vmc
import copy
import numpy as np
print("running")
df, coords = vmc(copy.copy(wf),
np.asarray(lastrun[1]).copy(),
nsteps=nsteps,
accumulators=copy.copy(accumulators),
stepoffset=stepoffset)
return df, coords.tolist()
def test_accumulator(C2_ccecp_rhf):
"""Tests that the accumulator gets inserted into the data output correctly."""
mol, mf = C2_ccecp_rhf
nconf = 500
wf = Slater(mol, mf)
coords = initial_guess(mol, nconf)
df, coords = vmc(wf,
coords,
nsteps=30,
accumulators={"energy": EnergyAccumulator(mol)})
df = pd.DataFrame(df)
eaccum = EnergyAccumulator(mol)
eaccum_energy = eaccum(coords, wf)
df = pd.DataFrame(df)
print(df["energyke"][29] == np.average(eaccum_energy["ke"]))
assert df["energyke"][29] == np.average(eaccum_energy["ke"])
def test_accumulator():
"""Tests that the accumulator gets inserted into the data output correctly."""
mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5",
basis="cc-pvtz",
unit="bohr",
verbose=5)
mf = scf.RHF(mol).run()
nconf = 5000
wf = Slater(mol, mf)
coords = initial_guess(mol, nconf)
df, coords = vmc(wf,
coords,
nsteps=30,
accumulators={"energy": EnergyAccumulator(mol)})
df = pd.DataFrame(df)
eaccum = EnergyAccumulator(mol)
eaccum_energy = eaccum(coords, wf)
df = pd.DataFrame(df)
print(df["energytotal"][29] == np.average(eaccum_energy["total"]))
assert df["energytotal"][29] == np.average(eaccum_energy["total"])
def test_vmc(C2_ccecp_rhf):
"""
Test that a VMC calculation of a Slater determinant matches Hartree-Fock within error bars.
"""
mol, mf = C2_ccecp_rhf
nconf = 500
nsteps = 300
warmup = 30
wf = Slater(mol, mf)
coords = initial_guess(mol, nconf)
df, coords = vmc(
wf,
coords,
nblocks=int(nsteps / 30),
nsteps_per_block=30,
accumulators={"energy": EnergyAccumulator(mol)},
)
df = pd.DataFrame(df)["energytotal"][int(warmup / 30):]
en = df.mean()
err = df.sem()
assert en - mf.energy_tot(
) < 5 * err, "pyscf {0}, vmc {1}, err {2}".format(mf.energy_tot(), en, err)
def test_ecp():
mol = gto.M(atom="C 0. 0. 0.", ecp="bfd", basis="bfd_vtz")
mf = scf.RHF(mol).run()
nconf = 5000
wf = Slater(mol, mf)
coords = initial_guess(mol, nconf)
df, coords = vmc(wf,
coords,
nsteps=100,
accumulators={"energy": EnergyAccumulator(mol)})
df = pd.DataFrame(df)
warmup = 30
print(
"mean field",
mf.energy_tot(),
"vmc estimation",
np.mean(df["energytotal"][warmup:]),
np.std(df["energytotal"][warmup:]),
)
assert abs(mf.energy_tot() -
np.mean(df["energytotal"][warmup:])) <= np.std(
df["energytotal"][warmup:])
def test(atom="He", total_spin=0, total_charge=0, scf_basis="sto-3g"):
mol = gto.M(
atom="%s 0. 0. 0.; %s 0. 0. 1.5" % (atom, atom),
basis=scf_basis,
unit="bohr",
verbose=4,
spin=total_spin,
charge=total_charge,
)
mf = scf.UHF(mol).run()
# Intrinsic Atomic Orbitals
iaos = make_separate_spin_iaos(
mol, mf, np.array([i for i in range(mol.natm)]), iao_basis="minao"
)
# iaos=make_combined_spin_iaos(mol,mf,np.array([i for i in range(mol.natm)]),iao_basis='minao')
# MOs in the IAO basis
mo = reps_combined_spin_iaos(
iaos,
mf,
np.einsum("i,j->ji", np.arange(mf.mo_coeff[0].shape[1]), np.array([1, 1])),
)
# Mean-field obdm in IAO basis
mfobdm = mf.make_rdm1(mo, mf.mo_occ)
# Mean-field tbdm in IAO basis
mftbdm = singledet_tbdm(mf, mfobdm)
### Test TBDM calculation.
# VMC params
nconf = 500
n_vmc_steps = 400
vmc_tstep = 0.3
vmc_warmup = 30
# TBDM params
tbdm_sweeps = 4
tbdm_tstep = 0.5
wf = PySCFSlater(mol, mf) # Single-Slater (no jastrow) wf
configs = initial_guess(mol, nconf)
energy = EnergyAccumulator(mol)
obdm_up = OBDMAccumulator(mol=mol, orb_coeff=iaos[0], nsweeps=tbdm_sweeps, spin=0)
obdm_down = OBDMAccumulator(mol=mol, orb_coeff=iaos[1], nsweeps=tbdm_sweeps, spin=1)
tbdm_upup = TBDMAccumulator(
mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(0, 0)
)
tbdm_updown = TBDMAccumulator(
mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(0, 1)
)
tbdm_downup = TBDMAccumulator(
mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(1, 0)
)
tbdm_downdown = TBDMAccumulator(
mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(1, 1)
)
print("VMC...")
df, coords = vmc(
wf,
configs,
nsteps=n_vmc_steps,
tstep=vmc_tstep,
accumulators={
"energy": energy,
"obdm_up": obdm_up,
"obdm_down": obdm_down,
"tbdm_upup": tbdm_upup,
"tbdm_updown": tbdm_updown,
"tbdm_downup": tbdm_downup,
"tbdm_downdown": tbdm_downdown,
},
verbose=True,
)
# Compares obdm from QMC and MF
obdm_est = {}
for k in ["obdm_up", "obdm_down"]:
avg_norm = np.mean(df[k + "norm"][vmc_warmup:], axis=0)
avg_obdm = np.mean(df[k + "value"][vmc_warmup:], axis=0)
obdm_est[k] = normalize_obdm(avg_obdm, avg_norm)
qmcobdm = np.array([obdm_est["obdm_up"], obdm_est["obdm_down"]])
print("\nComparing QMC and MF obdm:")
for s in [0, 1]:
# print('QMC obdm[%d]:\n'%s,qmcobdm[s])
# print('MF obdm[%d]:\n'%s,mfobdm[s])
print("diff[%d]:\n" % s, qmcobdm[s] - mfobdm[s])
# Compares tbdm from QMC and MF
avg_norm = {}
avg_tbdm = {}
tbdm_est = {}
for t in ["tbdm_upup", "tbdm_updown", "tbdm_downup", "tbdm_downdown"]:
for k in df.keys():
if k.startswith(t + "norm_"):
avg_norm[k.split("_")[-1]] = np.mean(df[k][vmc_warmup:], axis=0)
if k.startswith(t + "value"):
avg_tbdm[k.split("_")[-1]] = np.mean(df[k][vmc_warmup:], axis=0)
for k in avg_tbdm:
tbdm_est[k] = normalize_tbdm(
avg_tbdm[k].reshape(2, 2, 2, 2), avg_norm["a"], avg_norm["b"]
)
qmctbdm = np.array(
[
[tbdm_est["upupvalue"], tbdm_est["updownvalue"]],
[tbdm_est["downupvalue"], tbdm_est["downdownvalue"]],
]
)
print("\nComparing QMC and MF tbdm:")
for sa, sb in [[0, 0], [0, 1], [1, 0], [1, 1]]:
# print('QMC tbdm[%d,%d]:\n'%(sa,sb),qmctbdm[sa,sb])
# print('MF tbdm[%d,%d]:\n'%(sa,sb),mftbdm[sa,sb])
diff = qmctbdm[sa, sb] - mftbdm[sa, sb]
print("diff[%d,%d]:\n" % (sa, sb), diff)
assert np.max(np.abs(diff)) < 0.05
def test():
from pyscf import gto, scf, lo
from numpy.linalg import solve
from pyqmc.slater import PySCFSlaterRHF
from pyqmc.mc import initial_guess, vmc
from pyqmc.accumulators import EnergyAccumulator
from pandas import DataFrame
### Generate some basic objects.
# Simple Li2 run.
mol = gto.M(atom='Li 0. 0. 0.; Li 0. 0. 1.5',
basis='sto-3g',
unit='bohr',
verbose=0)
mf = scf.RHF(mol).run()
# Lowdin orthogonalized AO basis.
lowdin = lo.orth_ao(mol, 'lowdin')
# MOs in the Lowdin basis.
mo = solve(lowdin, mf.mo_coeff)
# make AO to localized orbital coefficients.
mfobdm = mf.make_rdm1(mo, mf.mo_occ)
#print(mfobdm.diagonal().round(2))
### Test one-body sampler.
#test_sample_onebody(mol,lowdin,mf,nsample=int(1e4))
#test_sample_onebody(mol,lowdin,mf,nsample=int(4e4))
#test_sample_onebody(mol,lowdin,mf,nsample=int(1e5))
### Test OBDM calculation.
nconf = 500
nsteps = 400
obdm_steps = 2
warmup = 15
wf = PySCFSlaterRHF(mol, mf)
configs = initial_guess(mol, nconf)
energy = EnergyAccumulator(mol)
obdm = OBDMAccumulator(mol=mol, orb_coeff=mf.mo_coeff, nstep=obdm_steps)
df, coords = vmc(wf,
configs,
nsteps=nsteps,
accumulators={
'energy': energy,
'obdm': obdm
})
df = DataFrame(df)
df['obdm'] = df[['obdmvalue','obdmnorm']]\
.apply(lambda x:normalize_obdm(x['obdmvalue'],x['obdmnorm']),axis=1)
print(df[['obdmvalue', 'obdmnorm',
'obdm']].applymap(lambda x: x.ravel()[0]))
avg_norm = np.array(df.loc[warmup:,
'obdmnorm'].values.tolist()).mean(axis=0)
avg_obdm = np.array(df.loc[warmup:, 'obdm'].values.tolist()).mean(axis=0)
std_obdm = np.array(
df.loc[warmup:, 'obdm'].values.tolist()).std(axis=0) / nsteps**0.5
print("Average norm(orb)", avg_norm)
print("Average OBDM(orb,orb)", avg_obdm.diagonal().round(3))
print("OBDM error (orb,orb)",
std_obdm.diagonal().round(
3)) # Note this needs reblocking to be accurate.
print("AO occupation", mfobdm[0, 0])
print('mean field', mf.energy_tot(), 'vmc estimation',
np.mean(df['energytotal'][warmup:]),
np.std(df['energytotal'][warmup:]))
def test_pbc():
from pyscf.pbc import gto, scf
from pyqmc import PySCFSlaterUHF, PySCFSlaterPBC
from pyqmc import slaterpbc
import scipy
lvecs = (np.ones((3, 3)) - 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)) - 2 * np.eye(3)
mol = slaterpbc.get_supercell(mol, S)
kpts = slaterpbc.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 = PySCFSlaterPBC(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,
)
df = DataFrame(df)
obdm_est = {}
for k in ["obdm", "obdm_up", "obdm_down"]:
avg_norm = np.array(df.loc[warmup:, k + "norm"].values.tolist()).mean(axis=0)
avg_obdm = np.array(df.loc[warmup:, k + "value"].values.tolist()).mean(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
)
def rundmc(
wf,
configs,
weights=None,
tstep=0.01,
nsteps=1000,
branchtime=5,
stepoffset=0,
branchcut_start=3,
branchcut_stop=6,
drift_limiter=limdrift,
verbose=False,
accumulators=None,
ekey=("energy", "total"),
propagate=dmc_propagate,
feedback=1.0,
hdf_file=None,
client=None,
npartitions=None,
**kwargs,
):
"""
Run DMC
Args:
wf: A Wave function-like class. recompute(), gradient(), and updateinternals() are used, as well as anything (such as laplacian() ) used by accumulators
configs: (nconfig, nelec, 3) - initial coordinates to start calculation.
weights: (nconfig,) - initial weights to start calculation, defaults to uniform.
nsteps: number of DMC steps to take
tstep: Time step for move proposals. Introduces time step error.
branchtime: number of steps to take between branching
accumulators: A dictionary of functor objects that take in (coords,wf) and return a dictionary of quantities to be averaged. np.mean(quantity,axis=0) should give the average over configurations. If none, a default energy accumulator will be used.
ekey: tuple of strings; energy is needed for DMC weights. Access total energy by accumulators[ekey[0]](configs, wf)[ekey[1]
verbose: Print out step information
drift_limiter: a function that takes a gradient and a cutoff and returns an adjusted gradient
stepoffset: If continuing a run, what to start the step numbering at.
Returns: (df,coords,weights)
df: A list of dictionaries nstep long that contains all results from the accumulators.
coords: The final coordinates from this calculation.
weights: The final weights from this calculation
"""
# Restart from HDF file
if hdf_file is not None and os.path.isfile(hdf_file):
with h5py.File(hdf_file, "r") as hdf:
stepoffset = hdf["step"][-1] + 1
configs.load_hdf(hdf)
weights = np.array(hdf["weights"])
eref = hdf["eref"][-1]
esigma = hdf["esigma"][-1]
if verbose:
print("Restarted calculation")
else:
warmup = 2
df, configs = mc.vmc(
wf,
configs,
accumulators=accumulators,
client=client,
npartitions=npartitions,
verbose=verbose,
)
en = df[ekey[0] + ekey[1]][warmup:]
eref = np.mean(en).real
esigma = np.sqrt(np.var(en) * np.mean(df["nconfig"]))
if verbose:
print("eref start", eref, "esigma", esigma)
nconfig = configs.configs.shape[0]
if weights is None:
weights = np.ones(nconfig)
npropagate = int(np.ceil(nsteps / branchtime))
df = []
for step in range(npropagate):
if client is None:
df_, configs, weights = dmc_propagate(
wf,
configs,
weights,
tstep,
branchcut_start * esigma,
branchcut_stop * esigma,
eref=eref,
nsteps=branchtime,
accumulators=accumulators,
ekey=ekey,
drift_limiter=drift_limiter,
**kwargs,
)
else:
df_, configs, weights = dmc_propagate_parallel(
wf,
configs,
weights,
client,
npartitions,
tstep,
branchcut_start * esigma,
branchcut_stop * esigma,
eref=eref,
nsteps=branchtime,
accumulators=accumulators,
ekey=ekey,
drift_limiter=drift_limiter,
**kwargs,
)
df_["eref"] = eref
df_["step"] = step + stepoffset
df_["esigma"] = esigma
df_["tstep"] = tstep
df_["weight_std"] = np.std(weights)
df_["nsteps"] = branchtime
dmc_file(hdf_file, df_, {}, configs, weights)
# print(df_)
df.append(df_)
eref = (df_[ekey[0] + ekey[1]] - feedback * np.log(np.mean(weights))).real
configs, weights = branch(configs, weights)
if verbose:
print(
"energy",
df_[ekey[0] + ekey[1]],
"eref",
df_["eref"],
"sigma(w)",
df_["weight_std"],
)
df_ret = {}
for k in df[0].keys():
df_ret[k] = np.asarray([d[k] for d in df])
return df_ret, configs, weights
def test():
"""
Tests that the multi-slater wave function value, gradient and
parameter gradient evaluations are working correctly. Also
checks that VMC energy matches energy calculated in PySCF
"""
mol = gto.M(atom="Li 0. 0. 0.; H 0. 0. 1.5",
basis="cc-pvtz",
unit="bohr",
spin=0)
epsilon = 1e-4
delta = 1e-5
nsteps = 200
warmup = 10
for mf in [scf.RHF(mol).run(), scf.ROHF(mol).run(), scf.UHF(mol).run()]:
# Test same number of elecs
mc = mcscf.CASCI(mf, ncas=4, nelecas=(1, 1))
mc.kernel()
wf = MultiSlater(mol, mf, mc)
nconf = 10
nelec = np.sum(mol.nelec)
epos = initial_guess(mol, nconf)
for k, item in testwf.test_updateinternals(wf, epos).items():
assert item < epsilon
assert testwf.test_wf_gradient(wf, epos, delta=delta)[0] < epsilon
assert testwf.test_wf_laplacian(wf, epos, delta=delta)[0] < epsilon
assert testwf.test_wf_pgradient(wf, epos, delta=delta)[0] < epsilon
# Test same number of elecs
mc = mcscf.CASCI(mf, ncas=4, nelecas=(1, 1))
mc.kernel()
wf = pyqmc.default_msj(mol, mf, mc)[0]
nelec = np.sum(mol.nelec)
epos = initial_guess(mol, nconf)
for k, item in testwf.test_updateinternals(wf, epos).items():
assert item < epsilon
assert testwf.test_wf_gradient(wf, epos, delta=delta)[0] < epsilon
assert testwf.test_wf_laplacian(wf, epos, delta=delta)[0] < epsilon
assert testwf.test_wf_pgradient(wf, epos, delta=delta)[0] < epsilon
# Test different number of elecs
mc = mcscf.CASCI(mf, ncas=4, nelecas=(2, 0))
mc.kernel()
wf = MultiSlater(mol, mf, mc)
nelec = np.sum(mol.nelec)
epos = initial_guess(mol, nconf)
for k, item in testwf.test_updateinternals(wf, epos).items():
assert item < epsilon
assert testwf.test_wf_gradient(wf, epos, delta=delta)[0] < epsilon
assert testwf.test_wf_laplacian(wf, epos, delta=delta)[0] < epsilon
assert testwf.test_wf_pgradient(wf, epos, delta=delta)[0] < epsilon
# Quick VMC test
nconf = 1000
coords = initial_guess(mol, nconf)
df, coords = vmc(wf,
coords,
nsteps=nsteps,
accumulators={"energy": EnergyAccumulator(mol)})
df = pd.DataFrame(df)
df = reblock(df["energytotal"][warmup:], 20)
en = df.mean()
err = df.sem()
assert en - mc.e_tot < 5 * err
def rundmc(
wf,
configs,
weights=None,
tstep=0.01,
nsteps=1000,
branchtime=5,
stepoffset=0,
branchcut_start=10,
verbose=False,
accumulators=None,
ekey=("energy", "total"),
feedback=1.0,
hdf_file=None,
continue_from=None,
client=None,
npartitions=None,
**kwargs,
):
"""
Run DMC
:parameter wf: A Wave function-like class. recompute(), gradient(), and updateinternals() are used, as well as anything (such as laplacian() ) used by accumulators
:parameter configs: (nconfig, nelec, 3) - initial coordinates to start calculation.
:parameter weights: (nconfig,) - initial weights to start calculation, defaults to uniform.
:parameter nsteps: number of DMC steps to take
:parameter tstep: Time step for move proposals. Introduces time step error.
:parameter branchtime: number of steps to take between branching
:parameter accumulators: A dictionary of functor objects that take in (coords,wf) and return a dictionary of quantities to be averaged. np.mean(quantity,axis=0) should give the average over configurations. If none, a default energy accumulator will be used.
:parameter ekey: tuple of strings; energy is needed for DMC weights. Access total energy by accumulators[ekey[0]](configs, wf)[ekey[1]
:parameter verbose: Print out step information
:parameter stepoffset: If continuing a run, what to start the step numbering at.
:returns: (df,coords,weights)
df: A list of dictionaries nstep long that contains all results from the accumulators.
coords: The final coordinates from this calculation.
weights: The final weights from this calculation
"""
# Don't continue onto a file that's already there.
if continue_from is not None and hdf_file is not None and os.path.isfile(
hdf_file):
raise RuntimeError(
f"continue_from is set but hdf_file={hdf_file} already exists! Delete or rename {hdf_file} and try again."
)
# Restart if hdf_file is there
if continue_from is None and hdf_file is not None and os.path.isfile(
hdf_file):
continue_from = hdf_file
# Now we should be sure that there is a file
# to continue from, if given.
if continue_from is not None:
with h5py.File(continue_from, "r") as hdf:
stepoffset = hdf["step"][-1] + 1
configs.load_hdf(hdf)
weights = np.array(hdf["weights"])
if "e_trial" not in hdf.keys():
raise ValueError(
"Did not find e_trial in the restart file. This may mean that you are trying to restart from a different version of DMC"
)
e_trial = hdf["e_trial"][-1]
e_est = hdf["e_est"][-1]
esigma = hdf["esigma"][-1]
if verbose:
print(
f"Restarting calculation {continue_from} from step {stepoffset}"
)
else:
warmup = 2
df, configs = mc.vmc(
wf,
configs,
accumulators={ekey[0]: accumulators[ekey[0]]},
client=client,
npartitions=npartitions,
verbose=verbose,
)
en = df[ekey[0] + ekey[1]][warmup:]
eref = np.mean(en).real
e_trial = eref
e_est = eref
esigma = np.sqrt(np.var(en) * np.mean(df["nconfig"]))
if verbose:
print("eref start", eref, "esigma", esigma)
nconfig = configs.configs.shape[0]
if weights is None:
weights = np.ones(nconfig)
npropagate = int(np.ceil(nsteps / branchtime))
df = []
for step in range(npropagate):
if client is None:
df_, configs, weights = dmc_propagate(
wf,
configs,
weights,
tstep,
branchcut_start * esigma,
e_trial=e_trial,
e_est=e_est,
nsteps=branchtime,
accumulators=accumulators,
ekey=ekey,
**kwargs,
)
else:
df_, configs, weights = dmc_propagate_parallel(
wf,
configs,
weights,
client,
npartitions,
tstep,
branchcut_start * esigma,
e_trial=e_trial,
e_est=e_est,
nsteps=branchtime,
accumulators=accumulators,
ekey=ekey,
**kwargs,
)
df_["e_trial"] = e_trial
df_["e_est"] = e_est
df_["step"] = step + stepoffset
df_["esigma"] = esigma
df_["tstep"] = tstep
df_["weight_std"] = np.std(weights)
df_["nsteps"] = branchtime
dmc_file(hdf_file, df_, {}, configs, weights)
df.append(df_)
e_est = estimate_energy(hdf_file, df, ekey)
e_trial = e_est - feedback * np.log(np.mean(weights)).real
configs, weights = branch(configs, weights)
if verbose:
print(
"energy",
df_[ekey[0] + ekey[1]],
"e_trial",
e_trial,
"e_est",
e_est,
"sigma(w)",
df_["weight_std"],
)
df_ret = {k: np.asarray([d[k] for d in df]) for k in df[0].keys()}
return df_ret, configs, weights
def optimize_orthogonal(
wfs,
coords,
pgrad,
Starget=None,
forcing=None,
tstep=0.1,
max_iterations=30,
warmup=1,
warmup_options=None,
Ntarget=0.5,
max_step=10.0,
hdf_file=None,
linemin=True,
step_offset=0,
Ntol=0.05,
weight_boundaries=0.3,
sample_options=None,
correlated_options=None,
client=None,
npartitions=None,
verbose=True,
):
r"""
Minimize
.. math:: f(p_f) = E_f + \sum_i \lambda_{i=0}^{f-1} |S_{fi} - S_{fi}^*|^2
Where
.. math:: N_i = \langle \Psi_i | \Psi_i \rangle
.. math:: S_{fi} = \frac{\langle \Psi_f | \Psi_i \rangle}{\sqrt{N_f N_i}}
The \*'d and lambda values are respectively targets and forcings. f is the final wave function in the wave function array.
We only optimize the parameters of the final wave function, so all 'p' values here represent a parameter in the final wave function.
**Important arguments**
:wfs: a list of wave function objects. The last one is optimized; the rest are kept fixed and used as orthogonalization references
:coords: A Coord set
:pgrad: A Pgradient object
:tstep: Maximum timestep for line minimization, or timestep when line minimization is off
:max_iterations: Number of optimization steps to take
:warmup_options: a dictionary of options for warm up vmc
:Starget: An array-like of length len(wfs)-1, which indicates the target overlap for each reference wave function.
:forcing: An array-like of length len(wfs)-1, which gives the penalty (lambda) for each reference wave function
:hdf_file: A string that gives the filename to save the optimization data in.
**Arguments for experts**
Other arguments should not be changed unless you know what you're doing.
**Details about the implementation**
The derivatives are:
.. math:: \partial_p N_f = 2 Re \langle \partial_p \Psi_f | \Psi_f \rangle
.. math:: \langle \partial_p \Psi_f | \Psi_i \rangle = \int{ \frac{ \Psi_i\partial_p \Psi_f^*}{\rho} \frac{\rho}{\int \rho} }
.. math:: \partial_p S_{fi} = \frac{\langle \partial_p \Psi_f | \Psi_i \rangle}{\sqrt{N_f N_i}} - \frac{\langle \Psi_f | \Psi_i \rangle}{2\sqrt{N_f N_i}} \frac{\partial_p N_f}{N_f}
In this implementation, we set
.. math:: \rho = \sum_i |\Psi_i|^2
Note that in the definition of N there is an arbitrary normalization of rho. The targets are set relative to the normalization of the reference wave functions.
Some implementation notes regarding the normalization:
It's important for the normalization of the wave functions to be similar; otherwise the weights of one dominate and only one of the wave functions gets sampled.
Some notes:
* One could modify the relative weights in the definition of rho, but it's more convenient to output wave functions that are normalized with respect to each other.
* Adding a penalty to the normalization turns out to be fairly unstable.
* Moves to reduce the overlap and energy tend to change the normalization a lot (keep in mind that both the determinant and Jastrow parts can have gauge degrees of freedom). This can lead into a tailspin effect pretty quickly.
In this implementation, we handle this using three techniques:
* Most importantly, the cost function derivative is orthogonalized to the derivative of the normalization.
* In the line minimization, if the normalization deviates too far from 0.5 relative to the reference wave function, we do not consider the move. The correlated sampling is unreliable in that situation anyway.
* The wave function is renormalized if its normalization deviates too far from 0.5 relative to the first wave function.
"""
# Restart
if hdf_file is not None and os.path.isfile(hdf_file):
with h5py.File(hdf_file, "r") as hdf:
if "wf" in hdf.keys():
grp = hdf["wf"]
for k in grp.keys():
wfs[-1].parameters[k] = np.array(grp[k])
if "iteration" in hdf.keys():
step_offset = np.max(hdf["iteration"][...]) + 1
parameters = pgrad.transform.serialize_parameters(wfs[-1].parameters)
if Starget is None:
Starget = np.zeros(len(wfs) - 1)
if forcing is None:
forcing = np.ones(len(wfs) - 1)
Starget = np.asarray(Starget)
forcing = np.asarray(forcing)
if len(forcing) != len(wfs) - 1:
raise AttributeError(
"forcing should be an array of length equal to the wfs minus 1: " +
str(len(forcing)))
attr = dict(
tstep=tstep,
max_iterations=max_iterations,
forcing=forcing,
warmup=warmup,
Starget=Starget,
Ntarget=Ntarget,
max_step=max_step,
)
conditioner = pyqmc.linemin.sd_update
if sample_options is None:
sample_options = {}
if correlated_options is None:
correlated_options = {}
if client is None:
sampler = sample_overlap
correlated_sampler = correlated_sample
else:
sampler = dist_sample_overlap
correlated_sampler = dist_correlated_sample
sample_options["client"] = client
sample_options["npartitions"] = npartitions
correlated_options["client"] = client
correlated_options["npartitions"] = npartitions
# warm up to equilibrate the configurations before running optimization
if warmup_options is None:
warmup_options = dict(nblocks=1, nsteps=10)
if "tstep" not in warmup_options and "tstep" in sample_options:
warmup_options["tstep"] = sample_options["tstep"]
data, coords = mc.vmc(wfs[-1],
coords,
accumulators={},
client=client,
npartitions=npartitions,
**warmup_options)
# One set of configurations for every wave function
allcoords = [coords.copy() for _ in wfs[:-1]]
dtype = np.complex if wfs[-1].iscomplex else np.float
for step in range(max_iterations):
# we iterate until the normalization is reasonable
# One could potentially save a little time here by not computing the gradients
# every time, but typically we don't have to renormalize if the moves are good
# Memory efficient implementation
nwf = len(wfs)
normalization = np.zeros(nwf - 1)
total_energy = 0
# energy_derivative = np.zeros(len(parameters))
N_derivative = np.zeros(len(parameters))
condition = np.zeros((len(parameters), len(parameters)))
overlaps = np.zeros(nwf - 1, dtype=dtype)
overlap_derivatives = np.zeros((nwf - 1, len(parameters)), dtype=dtype)
while True:
return_data, _ = sampler([wfs[0], wfs[-1]], allcoords[0], pgrad,
**sample_options)
tmp_deriv = evaluate(return_data, warmup)
N = tmp_deriv["N"][-1]
if verbose:
print("Normalization", N, flush=True)
if abs(N - Ntarget) < Ntol:
normalization[0] = tmp_deriv["N"][-1]
total_energy += tmp_deriv["total"] / (nwf - 1)
energy_derivative = tmp_deriv["energy_derivative"] / (nwf - 1)
N_derivative += tmp_deriv["N_derivative"] / (nwf - 1)
condition += tmp_deriv["condition"] / (nwf - 1)
overlaps[0] = tmp_deriv["S"][-1, 0]
overlap_derivatives[0] = tmp_deriv["S_derivative"][0, :]
break
else:
renormalize([wfs[0], wfs[-1]], N)
parameters = pgrad.transform.serialize_parameters(
wfs[-1].parameters)
for i, wf in enumerate(wfs[1:-1]):
return_data, _ = sampler([wf, wfs[-1]], allcoords[i + 1], pgrad,
**sample_options)
deriv_data = evaluate(return_data, warmup)
normalization[i + 1] = deriv_data["N"][-1]
total_energy += deriv_data["total"] / (nwf - 1)
energy_derivative += deriv_data["energy_derivative"] / (nwf - 1)
N_derivative += deriv_data["N_derivative"] / (nwf - 1)
condition += deriv_data["condition"] / (nwf - 1)
overlaps[i + 1] = deriv_data["S"][-1, 0]
overlap_derivatives[i + 1] = deriv_data["S_derivative"][0, :]
if verbose:
print("normalization", normalization)
delta = overlaps - Starget
delta_phase = delta / np.abs(delta)
overlap_derivative = np.einsum(
"j,jk->k",
2.0 * forcing * np.abs(delta),
np.real(overlap_derivatives / delta_phase[:, np.newaxis]),
)
total_derivative = energy_derivative + overlap_derivative
if verbose:
print("############################# iteration ", step)
format_str = "{:<15}" * 2 + "{:<20.3}" * 2
print(format_str.format("Quantity", "wf", "val", "|g|"))
print(
format_str.format(
"energy",
len(wfs) - 1,
total_energy,
np.linalg.norm(energy_derivative),
))
print(
format_str.format("norm",
len(wfs) - 1, N,
np.linalg.norm(N_derivative)))
for i in range(len(wfs) - 1):
print(
format_str.format(
"overlap",
i,
overlaps[i],
np.linalg.norm(overlap_derivatives[i]),
),
flush=True,
)
# Use SR to condition the derivatives
total_derivative, N_derivative = np.einsum(
"ij,dj->di",
np.linalg.inv(condition + 0.1 * np.eye(condition.shape[0])),
[total_derivative, N_derivative],
)
# Try to move in the projection that doesn't change the norm
# Here we project out the norm derivative
if np.linalg.norm(N_derivative) > 1e-8:
total_derivative -= (np.dot(total_derivative, N_derivative) *
N_derivative /
(np.linalg.norm(N_derivative))**2)
deriv_norm = np.linalg.norm(total_derivative)
if deriv_norm > max_step:
total_derivative = total_derivative * max_step / deriv_norm
test_tsteps = np.linspace(-0.1 * tstep, tstep, 11)
test_parameters = [
parameters + conditioner(total_derivative, condition, x)
for x in test_tsteps
]
data = []
for icoord, wf in zip(allcoords, wfs):
data.append(
correlated_sampler([wf, wfs[-1]], icoord, test_parameters,
pgrad, **correlated_options))
line_data = {}
for k in data[0].keys():
line_data[k] = np.asarray([x[k] for x in data])
yfit = []
xfit = []
overlap_cost = (
forcing[:, np.newaxis] *
np.abs(line_data["overlap"][:, :, 0] - Starget[:, np.newaxis])**2)
cost = np.mean(line_data["total"], axis=0) + np.sum(overlap_cost,
axis=0)
mask = (np.abs(line_data["weight"] - 1.0) > weight_boundaries) & (
np.abs(line_data["weight"]) > weight_boundaries)
mask = np.all(mask, axis=0)
if verbose:
print("tsteps", test_tsteps)
print("cost", cost)
print("overlap cost", overlap_cost)
print("mask", mask)
xfit = test_tsteps[mask]
yfit = cost[mask]
if verbose:
print("|total_derivative|", np.linalg.norm(total_derivative))
if len(xfit) > 2:
min_tstep = pyqmc.linemin.stable_fit(xfit, yfit)
if verbose:
print("chose to move", min_tstep, flush=True)
parameters = parameters + conditioner(total_derivative, condition,
min_tstep)
else:
print("WARNING: did not find valid moves. Reducing the timestep")
tstep *= 0.5
for k, it in pgrad.transform.deserialize(parameters).items():
wfs[-1].parameters[k] = it
save_data = {
"energy": total_energy,
"overlap": overlaps,
"gradient": total_derivative,
"N": N,
"parameters": parameters,
"iteration": step + step_offset,
"normalization": normalization,
"overlap_derivatives": overlap_derivatives,
"energy_derivative": energy_derivative,
"line_tsteps": test_tsteps,
"line_cost": cost,
"line_norm": line_data["weight"],
}
ortho_hdf(hdf_file, save_data, attr, coords, wfs[-1].parameters)
return wfs
