def dpH2(x, pgrad, node_coords, gradient, wf):
"""Variance dpH^2 across a node"""
coords = OpenConfigs(gradient * x + node_coords)
val = wf.recompute(coords)
d = pgrad(coords, wf)
dpH2 = np.array(d['dpH'])[:, 0, 0]**2 * np.exp(2 * val[1])
return dpH2
def test():
# Default multi-Slater wave function
mol = gto.M(atom="Li 0. 0. 0.; H 0. 0. 1.5",
basis="cc-pvtz",
unit="bohr",
spin=0)
mf = scf.RHF(mol).run()
mc = mcscf.CASCI(mf, ncas=2, nelecas=(1, 1))
mc.kernel()
wf, to_opt = default_msj(mol, mf, mc)
old_parms = wf.parameters
lt = LinearTransform(wf.parameters, to_opt)
# Test serialize parameters
x0 = lt.serialize_parameters(wf.parameters)
x0 += np.random.normal(size=x0.shape)
wf.parameters = lt.deserialize(x0)
assert wf.parameters["wf1det_coeff"][0] == old_parms["wf1det_coeff"][0]
assert np.sum(wf.parameters["wf2bcoeff"][0] -
old_parms["wf2bcoeff"][0]) == 0
# Test serialize gradients
configs = OpenConfigs(np.random.randn(10, 4, 3))
wf.recompute(configs)
pgrad = wf.pgradient()
pgrad_serial = lt.serialize_gradients(pgrad)
assert np.sum(pgrad_serial[:, :3] - pgrad["wf1det_coeff"][:, 1:4]) == 0
def locatenode(df, scale):
"""
Pinpoint a node, scale variable scales wf value
so that minimization algorithm can converge
"""
from scipy.optimize import minimize_scalar
mol, wf, to_opt, freeze = wavefunction()
ind = np.argsort(-df['dpH'].values)
node_coords_0 = np.array(df.iloc[ind[0]]['configs']) #Pretty close!
#Get the gradient
def wfgrad(coords, wf, mol):
nelec = mol.nelec[0] + mol.nelec[1]
val = wf.recompute(coords)
grad = []
for e in range(nelec):
node_grad = wf.gradient(e, coords.electron(e)) * np.exp(val[1]) * val[0]
grad.append(node_grad)
grad = np.array(grad)
grad /= np.linalg.norm(grad.ravel())
return np.rollaxis(grad, -1)
#Value function along that gradient
def wvfal(x, wf, gradient):
node_coords = OpenConfigs(node_coords_0 + gradient * x)
val = wf.recompute(node_coords)
return np.exp(2 * val[1])/scale #Scaling for minimization
#Minimize function
node_coords = node_coords_0
for i in range(1):
val = wf.recompute(OpenConfigs(node_coords))
grad = wfgrad(OpenConfigs(node_coords), wf, mol)
print("Wfval: ", np.exp(val[1])*val[0])
res = minimize_scalar(lambda x: wvfal(x, wf, grad), bracket = [-0.1, 0.1], tol = 1e-16)
print("x: ",res.x)
#Upgrade gradient
node_coords += grad * res.x[0]
val = wf.recompute(OpenConfigs(node_coords))
print("Wfval post: ", np.exp(val[1])*val[0])
return node_coords, grad
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 initial_guess(mol, nconfig, r=1.0):
""" Generate an initial guess by distributing electrons near atoms
proportional to their charge.
Args:
mol: A PySCF-like molecule object. Should have atom_charges(), atom_coords(), and nelec
nconfig: How many configurations to generate.
r: How far from the atoms to distribute the electrons
Returns:
A numpy array with shape (nconfig,nelectrons,3) with the electrons randomly distributed near
the atoms.
"""
from pyqmc.coord import OpenConfigs, PeriodicConfigs
nelec = np.sum(mol.nelec)
epos = np.zeros((nconfig, nelec, 3))
wts = mol.atom_charges()
wts = wts / np.sum(wts)
# assign electrons to atoms based on atom charges
# assign the minimum number first, and assign the leftover ones randomly
# this algorithm chooses atoms *with replacement* to assign leftover electrons
for s in [0, 1]:
neach = np.array(np.floor(mol.nelec[s] * wts),
dtype=int) # integer number of elec on each atom
nleft = (mol.nelec[s] * wts - neach
) # fraction of electron unassigned on each atom
nassigned = np.sum(neach) # number of electrons assigned
totleft = int(mol.nelec[s] -
nassigned) # number of electrons not yet assigned
if totleft > 0:
bins = np.cumsum(nleft) / totleft
inds = np.argpartition(np.random.random((nconfig, len(wts))),
totleft,
axis=1)[:, :totleft]
ind0 = s * mol.nelec[0]
epos[:, ind0:ind0 + nassigned, :] = np.repeat(
mol.atom_coords(), neach,
axis=0)[np.newaxis] # assign core electrons
epos[:,
ind0 + nassigned:ind0 + mol.nelec[s], :] = mol.atom_coords()[
inds] # assign remaining electrons
epos += r * np.random.randn(
*epos.shape) # random shifts from atom positions
if hasattr(mol, "a"):
epos = PeriodicConfigs(epos, mol.a)
else:
epos = OpenConfigs(epos)
return epos
def dpH(x, pgrad, pgrad_bare, node_coords, gradient, wf):
"""Bias of dpH across a node"""
coords = OpenConfigs(gradient * x + node_coords)
val = wf.recompute(coords)
d = pgrad(coords, wf)
dpH = np.array(d['dpH'])[:, 0, 0]
d_bare = pgrad_bare(coords, wf)
dpH_bare = np.array(d_bare['dpH'])[:, 0]
return (dpH - dpH_bare) * np.exp(2 * val[1])
def viznode(node_coords, node_grad, cutoffs, vizfile='viznode.pdf'):
from wavefunction import wavefunction
mol, wf, to_opt, freeze = wavefunction()
eacc = EnergyAccumulator(mol)
transform = LinearTransform(wf.parameters, to_opt, freeze)
#Generate the distances we want
x = np.arange(np.sqrt(1 / 0.01), 1001)
x = 1. / x**2
x = np.append(x, np.linspace(0.01, 0.12, 100))
x = np.append(x, np.array([0, 1e-15, 1e-12, 1e-8] + [-y for y in x]))
x = np.sort(x)
coord_path = np.einsum('ij,l->lij', node_grad[0], x) + node_coords[0]
coord_path = OpenConfigs(coord_path)
#Move across the node in this path
val = wf.recompute(coord_path)
fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(3, 6), sharex=True)
for k, cutoff in enumerate([1e-8] + cutoffs):
pgrad = PGradTransform_new(eacc,
transform,
nodal_cutoff=np.array([cutoff]))
d = pgrad(coord_path, wf)
total = d['total']
dpH = np.array(d['dpH'])[:, 0, 0]
if (cutoff == 1e-8):
ax[0].plot(x,
dpH * (val[0] * np.exp(val[1]))**2,
'k-',
label=r'$10^{' + str(int(np.log10(cutoff))) + '}$')
ax[1].plot(x, np.log10(dpH**2 * (val[0] * np.exp(val[1]))**2),
'k-')
else:
ax[0].plot(x,
dpH * (val[0] * np.exp(val[1]))**2,
'-',
label=r'$10^{' + str(int(np.log10(cutoff))) + '}$')
ax[1].plot(x, np.log10(dpH**2 * (val[0] * np.exp(val[1]))**2), '-')
ax[0].set_ylabel(r'$E_L\frac{\partial_p \Psi}{\Psi} f_\epsilon |\Psi|^2$')
ax[1].set_ylabel(
r'log$_{10}((E_L\frac{\partial_p \Psi}{\Psi})^2 f_\epsilon^2|\Psi|^2)$'
)
ax[1].set_xlabel(r'$l$ (Bohr)')
ax[0].set_xlim((-max(x) - 0.02, max(x) + 0.02))
ax[1].set_xlim((-max(x) - 0.02, max(x) + 0.02))
ax[0].legend(loc='best', title=r'$\epsilon$ (Bohr)')
plt.savefig(vizfile, bbox_inches='tight')
plt.close()
def collectconfigs(n, dump_file):
"""
Collect all the configurations from genconfig
into a single place
"""
dpH_total = []
weight_total = []
logweight_total = []
distance_squared = []
mol, wf, to_opt, freeze = wavefunction()
eacc = EnergyAccumulator(mol)
transform = LinearTransform(wf.parameters, to_opt, freeze)
pgrad_bare = PGradTransform_new(eacc, transform, 1e-20)
for i in range(1, n + 1):
print(i)
coords = OpenConfigs(pd.read_pickle('vmc/coords' + str(i) + '.pickle'))
df = pd.read_json('vmc/evals' + str(i) + '.json')
wf.recompute(coords)
print("Recomputed")
node_cut, r2 = pgrad_bare._node_cut(coords, wf)
print("nodes cut")
dpH = df['dpH'].values
wfval = df['wfval'].values
wfpval = df['wfpval'].values
logweight = 2 * (wfval - wfpval)
weight = np.exp(logweight)
print(i, weight[weight == 0])
dpH_total += list(dpH)
weight_total += list(weight)
logweight_total += list(logweight)
distance_squared += list(r2)
print(i, np.array(weight_total)[np.array(weight_total) == 0])
df = pd.DataFrame({
'dpH': dpH_total,
'weight_total': weight_total,
'logweight_total': logweight_total,
'distance_squared': distance_squared
})
df.to_pickle(dump_file)
return df
def sweepelectron():
"""
Sweep an electron across the molecule
to find a guess for the nodal position
"""
import copy
from pyqmc.accumulators import EnergyAccumulator, LinearTransform, PGradTransform
#Generate wave function and bare parameter gradient objects
mol, wf, to_opt, freeze = wavefunction()
eacc = EnergyAccumulator(mol)
transform = LinearTransform(wf.parameters, to_opt, freeze)
pgrad = PGradTransform(eacc, transform, 1e-20)
#Initial coords
configs = pyqmc.initial_guess(mol, 1).configs[:,:,:]
#Sweep electron 0
full_df = None
e = 3 #electron
dim = 1 #Coordinate to vary
for i in np.linspace(0, 20, 200):
new_configs = copy.deepcopy(configs)
new_configs[:,e,dim] += i
shifted_configs = OpenConfigs(new_configs)
wfval = wf.recompute(shifted_configs)
d = pgrad(shifted_configs, wf)
small_df = pd.DataFrame({
'ke':[d['ke'][0]],
'total':[d['total'][0]],
'dppsi':[d['dppsi'][0][0]],
'dpH' :[d['dpH'][0][0]],
'wfval':[wfval[0][0]*np.exp(wfval[1][0])],
'ycoord': i,
'configs':[copy.deepcopy(new_configs)],
})
if(full_df is None): full_df = small_df
else: full_df = pd.concat((full_df, small_df), axis=0)
return full_df.reset_index()
def psi2(x, node_coords, gradient, wf):
"""Wavefunction squared across the node"""
coords = OpenConfigs(gradient * x + node_coords)
val = wf.recompute(coords)
return np.exp(val[1] * 2)
def wvfal(x, wf, gradient):
node_coords = OpenConfigs(node_coords_0 + gradient * x)
val = wf.recompute(node_coords)
return np.exp(2 * val[1])/scale #Scaling for minimization
import pyqmc
from pyqmc.slateruhf import PySCFSlaterUHF
# from pyqmc.jastrow import Jastrow2B
from pyqmc.coord import OpenConfigs
import time
import pandas as pd
mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5", basis="cc-pvtz", unit="bohr")
mf = scf.UHF(mol).run()
wf = PySCFSlaterUHF(mol, mf)
# wf=Jastrow2B(10,mol)
df = []
for i in range(5):
configs = OpenConfigs(np.random.randn(18000, np.sum(mol.nelec), 3))
res = test_wf_gradient_laplacian(wf, configs)
for d in res:
d.update({"step": i})
df.extend(res)
print("testing gradient: errors\n", pd.DataFrame(df))
quit()
for i in range(5):
configs = OpenConfigs(np.random.randn(10, np.sum(mol.nelec), 3))
print("testing gradient: errors",
test_wf_gradient(wf, configs, delta=1e-5))
for i in range(5):
configs = OpenConfigs(np.random.randn(10, np.sum(mol.nelec), 3))
print("testing laplacian: errors",
test_wf_laplacian(wf, configs, delta=1e-5))
import numpy as np from pyscf import lib, gto, scf import pyqmc from pyqmc.slateruhf import PySCFSlaterUHF from pyqmc.jastrowspin import JastrowSpin # from pyqmc.jastrow import Jastrow2B from pyqmc.coord import OpenConfigs import time from pyqmc.manybody_jastrow import J3 import pyqmc.testwf as test from pyqmc.wf import WaveFunction from pyqmc.multiplywf import MultiplyWF # import pandas as pd mol = gto.M(atom="Li 0. 0. 0.; Li 0. 0. 1.5", basis="sto-3g", unit="bohr") mf = scf.UHF(mol).run() wf1 = PySCFSlaterUHF(mol, mf) wf2 = JastrowSpin(mol) # wf3 = J3(mol) wf = WaveFunction([wf1, wf2]) wfmultiply = MultiplyWF(wf1, wf2) configs = OpenConfigs(np.random.randn(10, np.sum(mol.nelec), 3)) wf.recompute(configs) # res = test.test_wf_gradient(wf, configs) e = 2 epos = configs.electron(e) test.test_mask(wf, 3, epos)