import numpy as np
class Parameters:
def __init__(self, dicts):
self.data = {}
self.wf_count = len(dicts)
for (i, d) in enumerate(dicts):
self.data["wf" + str(i + 1)] = d
def __setitem__(self, idx, value):
k1 = idx[0:3]
k2 = idx[3:]
self.data[k1][k2] = value
def __getitem__(self, idx):
k1 = idx[0:3]
k2 = idx[3:]
return self.data[k1][k2]
def __delitem__(self, idx):
k1 = idx[0:3]
k2 = idx[3:]
del self.data[k1][k2]
def __iter__(self):
for i in range(self.wf_count):
k1 = "wf" + str(i + 1)
for k2 in self.data[k1].keys():
yield k1 + k2
def __len__(self):
return sum(len(i) for i in self.data)
def items(self):
for i in range(self.wf_count):
k1 = "wf" + str(i + 1)
for k2 in self.data[k1].keys():
yield k1 + k2, self.data[k1][k2]
def __repr__(self):
return "WFmerger: " + self.data.__repr__()
def keys(self):
for i in range(self.wf_count):
k1 = "wf" + str(i + 1)
for k2 in self.data[k1].keys():
yield k1 + k2
def values(self):
for i in range(self.wf_count):
k1 = "wf" + str(i + 1)
for k2 in self.data[k1].keys():
yield self.data[k1][k2]
[docs]class MultiplyWF:
"""
A general representation of a wavefunction as a product of multiple wf_factors
"""
def __init__(self, *wf_factors):
self.wf_factors = [*wf_factors]
self.parameters = Parameters([wf.parameters for wf in wf_factors])
self.iscomplex = bool(sum(wf.iscomplex for wf in wf_factors))
self.dtype = complex if self.iscomplex else float
def recompute(self, configs):
signs = np.ones(len(configs.configs))
vals = np.zeros(len(configs.configs))
for wf in self.wf_factors:
results = wf.recompute(configs)
signs = signs * results[0]
vals += results[1]
return signs, vals
def updateinternals(self, e, epos, mask=None):
for wf in self.wf_factors:
wf.updateinternals(e, epos, mask=mask)
def value(self):
results = [wf.value() for wf in self.wf_factors]
results = np.array([*results])
return np.prod(results[:, 0, :], axis=0), np.sum(results[:, 1, :], axis=0)
def gradient(self, e, epos):
grads = [wf.gradient(e, epos) for wf in self.wf_factors]
return np.sum(grads, axis=0)
def testvalue(self, e, epos, mask=None):
testvalues = [wf.testvalue(e, epos, mask=mask) for wf in self.wf_factors]
return np.prod(testvalues, axis=0)
def testvalue_many(self, e, epos, mask=None):
testvalues = [wf.testvalue_many(e, epos, mask=mask) for wf in self.wf_factors]
return np.prod(testvalues, axis=0)
def gradient_value(self, e, epos):
grad_vals = [wf.gradient_value(e, epos) for wf in self.wf_factors]
grads, vals = list(zip(*grad_vals))
return np.sum(grads, axis=0), np.prod(vals, axis=0)
def gradient_laplacian(self, e, epos):
grad_laps = [wf.gradient_laplacian(e, epos) for wf in self.wf_factors]
grads, laps = list(zip(*grad_laps))
cross_term = np.zeros(laps[0].shape, dtype=self.dtype)
nwf = len(self.wf_factors)
for i in range(nwf):
for j in range(i + 1, nwf):
cross_term += np.sum(grads[i] * grads[j], axis=0)
return np.sum(grads, axis=0), np.sum(laps, axis=0) + cross_term * 2
def laplacian(self, e, epos):
return self.gradient_laplacian(e, epos)[1]
def pgradient(self):
return Parameters([wf.pgradient() for wf in self.wf_factors])
def test_parameters():
dicts = [{"coeff" + str(i): np.random.rand(3)} for i in range(10)]
p = Parameters(dicts)
# test len
assert len(p) == 30
print("len test passed")
# test getitem
assert p["wf2coeff2"].all() == dicts[2]["coeff2"].all()
print("getitem test passed")
new_coeff = np.random.rand(5)
# test setitem
p["wf2coeff2"] = new_coeff
assert p["wf2coeff2"].all() == new_coeff.all()
print("setitem test passed")