def sample_overlap_worker(wfs, configs, pgrad, nsteps, tstep=0.5):
r"""Run nstep Metropolis steps to sample a distribution proportional to
:math:`\sum_i |\Psi_i|^2`, where :math:`\Psi_i` = wfs[i]
"""
nconf, nelec, _ = configs.configs.shape
for wf in wfs:
wf.recompute(configs)
block_avg = {}
block_avg["acceptance"] = np.zeros(nsteps)
for n in range(nsteps):
for e in range(nelec): # a sweep
# Propose move
grads = [
np.real(wf.gradient(e, configs.electron(e)).T) for wf in wfs
]
grad = mc.limdrift(np.mean(grads, axis=0))
gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3))
newcoorde = configs.configs[:, e, :] + gauss + grad * tstep
newcoorde = configs.make_irreducible(e, newcoorde)
# Compute reverse move
grads, vals = list(
zip(*[wf.gradient_value(e, newcoorde) for wf in wfs]))
grads = [np.real(g.T) for g in grads]
new_grad = mc.limdrift(np.mean(grads, axis=0))
forward = np.sum(gauss**2, axis=1)
backward = np.sum((gauss + tstep * (grad + new_grad))**2, axis=1)
# Acceptance
t_prob = np.exp(1 / (2 * tstep) * (forward - backward))
wf_ratios = np.abs(vals)**2
log_values = np.real(np.array([wf.value()[1] for wf in wfs]))
weights = np.exp(2 * (log_values - log_values[0]))
ratio = t_prob * np.sum(wf_ratios * weights,
axis=0) / weights.sum(axis=0)
accept = ratio > np.random.rand(nconf)
block_avg["acceptance"][n] += accept.mean() / nelec
# Update wave function
configs.move(e, newcoorde, accept)
for wf in wfs:
wf.updateinternals(e, newcoorde, configs, mask=accept)
# Collect rolling average
save_dat = collect_overlap_data(wfs, configs, pgrad)
for k, it in save_dat.items():
if k not in block_avg:
block_avg[k] = np.zeros((*it.shape, ), dtype=it.dtype)
if k in ["overlap", "overlap_gradient", "weight_final"]:
block_avg[k] += save_dat[k] / nsteps
else:
block_avg[k] += save_dat[k] * save_dat["weight_final"] / nsteps
for k, it in block_avg.items():
if k not in [
"overlap", "overlap_gradient", "weight_final", "acceptance"
]:
it /= block_avg["weight_final"]
return block_avg, configs
def limdrift_cutoff(g, tau, cutoff=1):
"""
Limit a vector to have a maximum magnitude of cutoff while maintaining direction
Args:
g: a [nconf,ndim] vector
cutoff: the maximum magnitude
Returns:
The vector with the cut off applied and multiplied by tau.
"""
return mc.limdrift(g, cutoff) * tau
def sample_overlap(wfs, configs, pgrad, nsteps=100, tstep=0.5):
r"""
Sample
.. math:: \rho(R) = \sum_i |\Psi_i(R)|^2
`pgrad` is expected to be a gradient generator. returns data as follows:
`overlap` :
.. math:: \left\langle \frac{\Psi_i^* \Psi_j}{\rho} \right\rangle
`overlap_gradient`:
.. math:: \left\langle \frac{\partial_{im} \Psi_i^* \Psi_j}{\rho} \right\rangle
Note that overlap_gradient is only saved for i = f, where f is the final wave function.
In addition, any key returned by `pgrad` will be saved for the final wave function.
"""
nconf, nelec, ndim = configs.configs.shape
for wf in wfs:
wf.recompute(configs)
return_data = {}
for step in range(nsteps):
# print("step", step)
for e in range(nelec):
# Propose move
grads = [
np.real(wf.gradient(e, configs.electron(e)).T) for wf in wfs
]
grad = limdrift(np.mean(grads, axis=0))
gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3))
newcoorde = configs.configs[:, e, :] + gauss + grad * tstep
newcoorde = configs.make_irreducible(e, newcoorde)
# Compute reverse move
grads = [np.real(wf.gradient(e, newcoorde).T) for wf in wfs]
new_grad = limdrift(np.mean(grads, axis=0))
forward = np.sum(gauss**2, axis=1)
backward = np.sum((gauss + tstep * (grad + new_grad))**2, axis=1)
# Acceptance
t_prob = np.exp(1 / (2 * tstep) * (forward - backward))
wf_ratios = np.array([wf.testvalue(e, newcoorde)**2 for wf in wfs])
log_values = np.array([wf.value()[1] for wf in wfs])
ref = log_values[0]
weights = np.exp(2 * (log_values - ref))
ratio = (t_prob * np.sum(wf_ratios * weights, axis=0) /
np.sum(weights, axis=0))
accept = ratio > np.random.rand(nconf)
# Update wave function
configs.move(e, newcoorde, accept)
for wf in wfs:
wf.updateinternals(e, newcoorde, mask=accept)
# print("accept", np.mean(accept))
log_values = np.array([wf.value() for wf in wfs])
# print(log_values.shape)
ref = np.max(log_values[:, 1, :], axis=0)
save_dat = {}
denominator = np.sum(np.exp(2 * (log_values[:, 1, :] - ref)), axis=0)
normalized_values = log_values[:, 0, :] * np.exp(log_values[:, 1, :] -
ref)
save_dat["overlap"] = np.mean(
np.einsum("ik,jk->ijk", normalized_values, normalized_values) /
denominator,
axis=-1,
)
weight = np.array([
np.exp(-2 * (log_values[i, 1, :] - log_values[:, 1, :]))
for i in range(len(wfs))
])
weight = 1.0 / np.sum(weight, axis=1)
dat = pgrad(configs, wfs[-1])
save_dat["overlap_gradient"] = np.mean(
np.einsum("km,k,jk->jmk", dat['dppsi'], normalized_values[-1],
normalized_values) / denominator,
axis=-1,
)
for k in dat.keys():
save_dat[k] = np.average(dat[k], axis=0, weights=weight[-1])
save_dat["weight"] = np.mean(weight, axis=1)
for k, it in save_dat.items():
if k not in return_data:
return_data[k] = np.zeros((nsteps, *it.shape))
return_data[k][step, ...] = it.copy()
return return_data, configs
def sample_overlap(wfs,
configs,
pgrad,
nblocks=10,
nsteps_per_block=10,
nsteps=None,
tstep=0.5):
r"""
Sample
.. math:: \rho(R) = \sum_i |\Psi_i(R)|^2
`pgrad` is expected to be a gradient generator. returns data as follows:
`overlap` :
.. math:: \left\langle \frac{\Psi_i^* \Psi_j}{\rho} \right\rangle
`overlap_gradient`:
.. math:: \left\langle \frac{\partial_{im} \Psi_i^* \Psi_j}{\rho} \right\rangle
Note that overlap_gradient is only saved for i = f, where f is the final wave function.
In addition, any key returned by `pgrad` will be saved for the final wave function.
"""
nconf, nelec, _ = configs.configs.shape
if nsteps is not None:
nblocks = nsteps
nsteps_per_block = 1
for wf in wfs:
wf.recompute(configs)
return_data = {}
for block in range(nblocks):
block_avg = {}
for step in range(nsteps_per_block):
for e in range(nelec):
# Propose move
grads = [
np.real(wf.gradient(e, configs.electron(e)).T)
for wf in wfs
]
grad = limdrift(np.mean(grads, axis=0))
gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3))
newcoorde = configs.configs[:, e, :] + gauss + grad * tstep
newcoorde = configs.make_irreducible(e, newcoorde)
# Compute reverse move
grads = [np.real(wf.gradient(e, newcoorde).T) for wf in wfs]
new_grad = limdrift(np.mean(grads, axis=0))
forward = np.sum(gauss**2, axis=1)
backward = np.sum((gauss + tstep * (grad + new_grad))**2,
axis=1)
# Acceptance
t_prob = np.exp(1 / (2 * tstep) * (forward - backward))
wf_ratios = np.array(
[wf.testvalue(e, newcoorde)**2 for wf in wfs])
log_values = np.array([wf.value()[1] for wf in wfs])
ref = log_values[0]
weights = np.exp(2 * (log_values - ref))
ratio = (t_prob * np.sum(wf_ratios * weights, axis=0) /
np.sum(weights, axis=0))
accept = ratio > np.random.rand(nconf)
# Update wave function
configs.move(e, newcoorde, accept)
for wf in wfs:
wf.updateinternals(e, newcoorde, mask=accept)
# print("accept", np.mean(accept))
log_values = np.array([wf.value() for wf in wfs])
# print(log_values.shape)
ref = np.max(log_values[:, 1, :], axis=0)
save_dat = {}
denominator = np.sum(np.exp(2 * (log_values[:, 1, :] - ref)),
axis=0)
normalized_values = log_values[:,
0, :] * np.exp(log_values[:, 1, :] -
ref)
save_dat["overlap"] = np.mean(
np.einsum("ik,jk->ijk", normalized_values, normalized_values) /
denominator,
axis=-1,
)
weight = np.array([
np.exp(-2 * (log_values[i, 1, :] - log_values[:, 1, :]))
for i in range(len(wfs))
])
weight = 1.0 / np.sum(weight, axis=1)
# Fast evaluation of dppsi_reg
dppsi = pgrad.transform.serialize_gradients(wfs[-1].pgradient())
node_cut, f = pgrad._node_regr(configs, wfs[-1])
dppsi_regularized = dppsi * f[:, np.newaxis]
save_dat["overlap_gradient"] = np.mean(
np.einsum(
"km,k,jk->jmk",
dppsi_regularized,
normalized_values[-1],
normalized_values,
) / denominator,
axis=-1,
)
# Weighted average on rest
dat = pgrad.avg(configs, wf, weight[-1])
for k in dat.keys():
save_dat[k] = dat[k]
save_dat["weight"] = np.mean(weight, axis=1)
# Rolling average within block
for k, it in save_dat.items():
if k not in block_avg:
block_avg[k] = np.zeros((*it.shape, ))
block_avg[k] += save_dat[k] / nsteps_per_block
# Blocks stored
for k, it in block_avg.items():
if k not in return_data:
return_data[k] = np.zeros((nblocks, *it.shape))
return_data[k][block, ...] = it.copy()
return return_data, configs
def sample_overlap_worker(wfs,
configs,
energy,
transforms,
nsteps=10,
nblocks=10,
tstep=0.5):
r"""Run nstep Metropolis steps to sample a distribution proportional to
:math:`\sum_i |\Psi_i|^2`, where :math:`\Psi_i` = wfs[i]
"""
nconf, nelec, _ = configs.configs.shape
for wf in wfs:
wf.recompute(configs)
weighted = []
unweighted = []
for block in range(nblocks):
weighted_block = {}
unweighted_block = {}
for n in range(nsteps):
for e in range(nelec): # a sweep
# Propose move
grads = [
np.real(wf.gradient(e, configs.electron(e)).T)
for wf in wfs
]
grad = mc.limdrift(np.mean(grads, axis=0))
gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3))
newcoorde = configs.configs[:, e, :] + gauss + grad * tstep
newcoorde = configs.make_irreducible(e, newcoorde)
# Compute reverse move
grads, vals = list(
zip(*[wf.gradient_value(e, newcoorde) for wf in wfs]))
grads = [np.real(g.T) for g in grads]
new_grad = mc.limdrift(np.mean(grads, axis=0))
forward = np.sum(gauss**2, axis=1)
backward = np.sum((gauss + tstep * (grad + new_grad))**2,
axis=1)
# Acceptance
t_prob = np.exp(1 / (2 * tstep) * (forward - backward))
wf_ratios = np.abs(vals)**2
log_values = np.real(np.array([wf.value()[1] for wf in wfs]))
weights = np.exp(2 * (log_values - log_values[0]))
ratio = t_prob * np.sum(wf_ratios * weights,
axis=0) / weights.sum(axis=0)
accept = ratio > np.random.rand(nconf)
#block_avg["acceptance"][n] += accept.mean() / nelec
# Update wave function
configs.move(e, newcoorde, accept)
for wf in wfs:
wf.updateinternals(e, newcoorde, configs, mask=accept)
# Collect rolling average
weighted_dat, unweighted_dat = collect_overlap_data(
wfs, configs, energy, transforms)
for k, it in unweighted_dat.items():
if k not in unweighted_block:
unweighted_block[k] = np.zeros((*it.shape, ),
dtype=it.dtype)
unweighted_block[k] += unweighted_dat[k] / nsteps
for k, it in weighted_dat.items():
if k not in weighted_block:
weighted_block[k] = [
np.zeros((*x.shape, ), dtype=x.dtype) for x in it
]
for b, v in zip(weighted_block[k], it):
b += v / nsteps
weighted.append(weighted_block)
unweighted.append(unweighted_block)
# here we modify the data so that it's a dictionary of lists of arrays for weighted
# and a dictionary of arrays for unweighted
# Access weighted as weighted[quantity][block, ...]
# Access unweighted as unweighted[quantity][block,...]
weighted = invert_list_of_dicts(weighted)
unweighted = invert_list_of_dicts(unweighted)
for k in weighted.keys():
weighted[k] = np.asarray(weighted[k])
for k in unweighted.keys():
unweighted[k] = np.asarray(unweighted[k])
return weighted, unweighted, configs
