class TestSlaterJastrowElectronCusp(unittest.TestCase):
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(atom='He 0.5 0 0; He -0.5 0 0',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = SlaterJastrow(mol,
jastrow_kernel=FullyConnectedJastrowKernel,
kinetic='jacobi',
include_all_mo=True,
configs='ground_state')
self.nbatch = 100
def test_ee_cusp(self):
import matplotlib.pyplot as plt
pos_x = torch.Tensor(np.random.rand(self.nbatch, self.wf.nelec, 3))
x = torch.linspace(0, 2, self.nbatch)
pos_x[:, 0, :] = torch.as_tensor([0., 0., 0.]) + 1E-6
pos_x[:, 1, 0] = 0.
pos_x[:, 1, 1] = 0.
pos_x[:, 1, 2] = x
pos_x[:, 2, :] = 0.5 * torch.as_tensor([1., 1., 1.])
pos_x[:, 3, :] = -0.5 * torch.as_tensor([1., 1., 1.])
pos_x = pos_x.reshape(self.nbatch, self.wf.nelec * 3)
pos_x.requires_grad = True
x = x.detach().numpy()
j = self.wf.jastrow(pos_x).detach().numpy()
plt.plot(x, j)
plt.show()
dx = x[1] - x[0]
dj = (j[1:] - j[0:-1]) / dx
plt.plot(x[:-1], dj / j[:-1])
plt.show()
epot = self.wf.electronic_potential(pos_x).detach().numpy()
ekin = self.wf.kinetic_energy_jacobi(pos_x).detach().numpy()
eloc = self.wf.local_energy(pos_x).detach().numpy()
plt.plot(x, epot)
plt.plot(x, ekin)
plt.plot(x, eloc)
plt.show()
class TestOrbitalWF(unittest.TestCase):
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(
atom='Li 0 0 0; H 0 0 1.',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = SlaterJastrow(mol,
kinetic='auto',
include_all_mo=True,
configs='cas(2,2)')
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.fc.weight.data = self.random_fc_weight
self.nbatch = 10
self.pos = torch.Tensor(
np.random.rand(self.nbatch, mol.nelec*3))
self.pos.requires_grad = True
def test_forward(self):
wfvals = self.wf(self.pos)
ref = torch.Tensor([[0.0522],
[0.0826],
[0.0774],
[0.1321],
[0.0459],
[0.0421],
[0.0551],
[0.0764],
[0.1164],
[0.2506]])
# assert torch.allclose(wfvals.data, ref, rtol=1E-4, atol=1E-4)
def test_antisymmetry(self):
"""Test that the wf values are antisymmetric
wrt exchange of 2 electrons of same spin."""
wfvals_ref = self.wf(self.pos)
if self.wf.nelec < 4:
print(
'Warning : antisymmetry cannot be tested with \
only %d electrons' % self.wf.nelec)
return
# test spin up
pos_xup = self.pos.clone()
perm_up = list(range(self.wf.nelec))
perm_up[0] = 1
perm_up[1] = 0
pos_xup = pos_xup.reshape(self.nbatch, self.wf.nelec, 3)
pos_xup = pos_xup[:, perm_up, :].reshape(
self.nbatch, self.wf.nelec*3)
wfvals_xup = self.wf(pos_xup)
assert(torch.allclose(wfvals_ref, -1*wfvals_xup))
# test spin down
pos_xdn = self.pos.clone()
perm_dn = list(range(self.wf.nelec))
perm_dn[self.wf.mol.nup-1] = self.wf.mol.nup
perm_dn[self.wf.mol.nup] = self.wf.mol.nup-1
pos_xdn = pos_xdn.reshape(self.nbatch, self.wf.nelec, 3)
pos_xdn = pos_xdn[:, perm_up, :].reshape(
self.nbatch, self.wf.nelec*3)
wfvals_xdn = self.wf(pos_xdn)
assert(torch.allclose(wfvals_ref, -1*wfvals_xdn))
def test_grad_mo(self):
"""Gradients of the MOs."""
mo = self.wf.pos2mo(self.pos)
dmo = self.wf.pos2mo(self.pos, derivative=1)
dmo_grad = grad(
mo,
self.pos,
grad_outputs=torch.ones_like(mo))[0]
gradcheck(self.wf.pos2mo, self.pos)
assert(torch.allclose(dmo.sum(), dmo_grad.sum()))
assert(torch.allclose(dmo.sum(-1),
dmo_grad.view(10, self.wf.nelec, 3).sum(-1)))
def test_hess_mo(self):
"""Hessian of the MOs."""
val = self.wf.pos2mo(self.pos)
d2val_grad = hess(val, self.pos)
d2val = self.wf.pos2mo(self.pos, derivative=2)
assert(torch.allclose(d2val.sum(), d2val_grad.sum()))
assert(torch.allclose(d2val.sum(-1).sum(-1),
d2val_grad.view(10, self.wf.nelec, 3).sum(-1).sum(-1)))
assert(torch.allclose(d2val.sum(-1),
d2val_grad.view(10, self.wf.nelec, 3).sum(-1)))
def test_local_energy(self):
self.wf.kinetic_energy = self.wf.kinetic_energy_autograd
eloc_auto = self.wf.local_energy(self.pos)
self.wf.kinetic_energy = self.wf.kinetic_energy_autograd
eloc_jac = self.wf.local_energy(self.pos)
ref = torch.Tensor([[-1.6567],
[-0.8790],
[-2.8136],
[-0.3644],
[-0.4477],
[-0.2709],
[-0.6964],
[-0.3993],
[-0.4777],
[-0.0579]])
# assert torch.allclose(
# eloc_auto.data, ref, rtol=1E-4, atol=1E-4)
assert torch.allclose(
eloc_auto.data, eloc_jac.data, rtol=1E-4, atol=1E-4)
def test_kinetic_energy(self):
eauto = self.wf.kinetic_energy_autograd(self.pos)
ejac = self.wf.kinetic_energy_jacobi(self.pos)
ref = torch.Tensor([[0.6099],
[0.6438],
[0.6313],
[2.0512],
[0.0838],
[0.2699],
[0.5190],
[0.3381],
[1.8489],
[5.2226]])
# assert torch.allclose(
# ejac.data, ref, rtol=1E-4, atol=1E-4)
assert torch.allclose(
eauto.data, ejac.data, rtol=1E-4, atol=1E-4)
def test_gradients_wf(self):
grads = self.wf.gradients_jacobi(self.pos)
grad_auto = self.wf.gradients_autograd(self.pos)
assert torch.allclose(grads, grad_auto)
def test_gradients_pdf(self):
grads_pdf = self.wf.gradients_jacobi(self.pos, pdf=True)
grads_auto = self.wf.gradients_autograd(self.pos, pdf=True)
assert torch.allclose(grads_pdf, grads_auto)
class TestGenericJastrowWF(unittest.TestCase):
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(
atom='Li 0 0 0; H 0 0 1.',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = SlaterJastrow(mol,
kinetic='auto',
configs='ground_state',
jastrow_kernel=FullyConnectedJastrowKernel)
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.fc.weight.data = self.random_fc_weight
self.nbatch = 10
self.pos = 1E-2 * torch.as_tensor(np.random.rand(
self.nbatch, self.wf.nelec*3))
self.pos.requires_grad = True
def test_forward(self):
wfvals = self.wf(self.pos)
def test_antisymmetry(self):
"""Test that the wf values are antisymmetric
wrt exchange of 2 electrons of same spin."""
wfvals_ref = self.wf(self.pos)
if self.wf.nelec < 4:
print(
'Warning : antisymmetry cannot be tested with \
only %d electrons' % self.wf.nelec)
return
# test spin up
pos_xup = self.pos.clone()
perm_up = list(range(self.wf.nelec))
perm_up[0] = 1
perm_up[1] = 0
pos_xup = pos_xup.reshape(self.nbatch, self.wf.nelec, 3)
pos_xup = pos_xup[:, perm_up, :].reshape(
self.nbatch, self.wf.nelec*3)
wfvals_xup = self.wf(pos_xup)
assert(torch.allclose(wfvals_ref, -1*wfvals_xup))
def test_grad_mo(self):
"""Gradients of the MOs."""
mo = self.wf.pos2mo(self.pos)
dmo = self.wf.pos2mo(self.pos, derivative=1)
dmo_grad = grad(
mo,
self.pos,
grad_outputs=torch.ones_like(mo))[0]
gradcheck(self.wf.pos2mo, self.pos)
assert(torch.allclose(dmo.sum(), dmo_grad.sum()))
assert(torch.allclose(dmo.sum(-1),
dmo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_hess_mo(self):
"""Hessian of the MOs."""
val = self.wf.pos2mo(self.pos)
d2val_grad = hess(val, self.pos)
d2val = self.wf.pos2mo(self.pos, derivative=2)
assert(torch.allclose(d2val.sum(), d2val_grad.sum()))
assert(torch.allclose(d2val.sum(-1).sum(-1),
d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1).sum(-1)))
assert(torch.allclose(d2val.sum(-1),
d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_local_energy(self):
self.wf.kinetic_energy = self.wf.kinetic_energy_autograd
eloc_auto = self.wf.local_energy(self.pos)
self.wf.kinetic_energy = self.wf.kinetic_energy_autograd
eloc_jac = self.wf.local_energy(self.pos)
assert torch.allclose(
eloc_auto.data, eloc_jac.data, rtol=1E-4, atol=1E-4)
def test_kinetic_energy(self):
eauto = self.wf.kinetic_energy_autograd(self.pos)
ejac = self.wf.kinetic_energy_jacobi(self.pos)
assert torch.allclose(
eauto.data, ejac.data, rtol=1E-4, atol=1E-4)
def test_gradients_wf(self):
grads = self.wf.gradients_jacobi(self.pos)
grad_auto = self.wf.gradients_autograd(self.pos)
assert torch.allclose(grads, grad_auto)
grads = grads.reshape(10, self.wf.nelec, 3)
grad_auto = grad_auto.reshape(10, self.wf.nelec, 3)
assert(torch.allclose(grads, grad_auto))
def test_gradients_pdf(self):
grads_pdf = self.wf.gradients_jacobi(self.pos, pdf=True)
grads_auto = self.wf.gradients_autograd(self.pos, pdf=True)
assert torch.allclose(grads_pdf, grads_auto)
class TestCompareSlaterJastrowOrbitalDependentBackFlow(unittest.TestCase):
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = SlaterJastrowBackFlow(mol,
kinetic='jacobi',
include_all_mo=True,
configs='single_double(2,2)',
backflow_kernel=BackFlowKernelInverse,
orbital_dependent_backflow=True)
for ker in self.wf.ao.backflow_trans.backflow_kernel.orbital_dependent_kernel:
ker.weight.data *= 0
self.wf_ref = SlaterJastrow(mol,
kinetic='jacobi',
include_all_mo=True,
configs='single_double(2,2)')
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.fc.weight.data = self.random_fc_weight
self.wf_ref.fc.weight.data = self.random_fc_weight
self.nbatch = 5
self.pos = torch.Tensor(np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
def test_forward(self):
"""Check that backflow give same results as normal SlaterJastrow."""
wf_val = self.wf(self.pos)
wf_val_ref = self.wf_ref(self.pos)
assert (torch.allclose(wf_val, wf_val_ref))
def test_jacobian_mo(self):
"""Check that backflow give same results as normal SlaterJastrow."""
dmo = self.wf.pos2mo(self.pos, derivative=1)
dmo_ref = self.wf_ref.pos2mo(self.pos, derivative=1)
assert (torch.allclose(dmo.sum(0), dmo_ref))
def test_hess_mo(self):
"""Check that backflow give same results as normal SlaterJastrow."""
d2ao = self.wf.ao(self.pos, derivative=2, sum_hess=False)
d2val = self.wf.ao2mo(d2ao)
d2ao_ref = self.wf_ref.ao(self.pos, derivative=2, sum_hess=True)
d2val_ref = self.wf_ref.ao2mo(d2ao_ref)
assert (torch.allclose(d2val_ref, d2val.sum(0)))
def test_grad_wf(self):
pass
def test_local_energy(self):
self.wf.kinetic_energy = self.wf.kinetic_energy_jacobi
eloc_jac = self.wf.local_energy(self.pos)
self.wf_ref.kinetic_energy = self.wf_ref.kinetic_energy_jacobi
eloc_jac_ref = self.wf_ref.local_energy(self.pos)
assert torch.allclose(eloc_jac_ref.data,
eloc_jac.data,
rtol=1E-4,
atol=1E-4)
def test_kinetic_energy(self):
ejac_ref = self.wf_ref.kinetic_energy_jacobi(self.pos)
ejac = self.wf.kinetic_energy_jacobi(self.pos)
assert torch.allclose(ejac_ref.data, ejac.data, rtol=1E-4, atol=1E-4)
