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 = CorrelatedOrbital(mol,
kinetic='auto',
configs='ground_state',
jastrow_type=FullyConnectedJastrow)
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, self.wf.nelec*3))
self.pos.requires_grad = True
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 = CorrelatedOrbital(mol,
kinetic='auto',
jastrow_type='pade_jastrow',
configs='single_double(2,4)',
include_all_mo=True)
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.fc.weight.data = self.random_fc_weight
self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.weight.shape)
self.nbatch = 3
self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
set_torch_double_precision()
# optimal parameters
self.opt_r = 0.69 # the two h are at +0.69 and -0.69
self.opt_sigma = 1.24
# molecule
self.mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# wave function
self.wf = CorrelatedOrbital(
self.mol,
kinetic='auto',
configs='cas(2,2)',
jastrow_type=FullyConnectedJastrow,
# jastrow_type='pade_jastrow',
include_all_mo=True)
# sampler
self.sampler = Metropolis(nwalkers=1000,
nstep=2000,
step_size=0.5,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'),
move={
'type': 'all-elec',
'proba': 'normal'
})
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# solver
self.solver = SolverOrbital(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# ground state energy
self.ground_state_energy = -1.16
# ground state pos
self.ground_state_pos = 0.69
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
set_torch_double_precision()
# molecule
path_hdf5 = (PATH_TEST / 'hdf5/LiH_adf_dz.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# wave function
self.wf = CorrelatedOrbital(self.mol,
kinetic='jacobi',
configs='cas(2,2)',
include_all_mo=True)
# fc weights
self.wf.fc.weight.data = torch.rand(self.wf.fc.weight.shape)
# jastrow weights
self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.weight.shape)
# sampler
self.sampler = Metropolis(nwalkers=500,
nstep=200,
step_size=0.05,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'),
move={
'type': 'all-elec',
'proba': 'normal'
})
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# solver
self.solver = SolverOrbital(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# artificial pos
self.nbatch = 10
self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
class TestLiHCorrelated(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
set_torch_double_precision()
# molecule
path_hdf5 = (PATH_TEST / 'hdf5/LiH_adf_dz.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# wave function
self.wf = CorrelatedOrbital(self.mol,
kinetic='jacobi',
configs='cas(2,2)',
include_all_mo=True)
# fc weights
self.wf.fc.weight.data = torch.rand(self.wf.fc.weight.shape)
# jastrow weights
self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.weight.shape)
# sampler
self.sampler = Metropolis(nwalkers=500,
nstep=200,
step_size=0.05,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'),
move={
'type': 'all-elec',
'proba': 'normal'
})
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# solver
self.solver = SolverOrbital(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# artificial pos
self.nbatch = 10
self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
def test_0_wavefunction(self):
eauto = self.wf.kinetic_energy_autograd(self.pos)
ejac = self.wf.kinetic_energy_jacobi(self.pos)
print(torch.stack([eauto, ejac], axis=1).squeeze())
assert torch.allclose(eauto.data, ejac.data, rtol=1E-4, atol=1E-4)
def test1_single_point(self):
# sample and compute observables
obs = self.solver.single_point()
e, v = obs.energy, obs.variance
# # values on different arch
# expected_energy = [-1.1464850902557373,
# -1.14937478612449]
# # values on different arch
# expected_variance = [0.9279592633247375,
# 0.7445300449383236]
# assert(np.any(np.isclose(e.data.item(), np.array(expected_energy))))
# assert(np.any(np.isclose(v.data.item(), np.array(expected_variance))))
def test2_wf_opt_grad_auto(self):
self.solver.sampler = self.sampler
self.solver.configure(track=['local_energy'],
loss='energy',
grad='auto')
obs = self.solver.run(5)
def test3_wf_opt_grad_manual(self):
self.solver.sampler = self.sampler
self.solver.configure(track=['local_energy'],
loss='energy',
grad='manual')
obs = self.solver.run(5)
class TestH2Correlated(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
set_torch_double_precision()
# optimal parameters
self.opt_r = 0.69 # the two h are at +0.69 and -0.69
self.opt_sigma = 1.24
# molecule
self.mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# wave function
self.wf = CorrelatedOrbital(
self.mol,
kinetic='auto',
configs='cas(2,2)',
jastrow_type=FullyConnectedJastrow,
# jastrow_type='pade_jastrow',
include_all_mo=True)
# sampler
self.sampler = Metropolis(nwalkers=1000,
nstep=2000,
step_size=0.5,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'),
move={
'type': 'all-elec',
'proba': 'normal'
})
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# solver
self.solver = SolverOrbital(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# ground state energy
self.ground_state_energy = -1.16
# ground state pos
self.ground_state_pos = 0.69
def test_0_wavefunction(self):
# artificial pos
self.nbatch = 10
self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
eauto = self.wf.kinetic_energy_autograd(self.pos)
ejac = self.wf.kinetic_energy_jacobi(self.pos)
print(torch.stack([eauto, ejac], axis=1).squeeze())
assert torch.allclose(eauto.data, ejac.data, rtol=1E-4, atol=1E-4)
def test1_single_point(self):
self.solver.wf.ao.atom_coords[0, 2] = -self.ground_state_pos
self.solver.wf.ao.atom_coords[1, 2] = self.ground_state_pos
self.solver.sampler = self.sampler
# sample and compute observables
obs = self.solver.single_point()
e, v = obs.energy, obs.variance
# values on different arch
# expected_energy = [-1.1286007165908813,
# -1.099538658544285]
# # values on different arch
# expected_variance = [0.45748308300971985,
# 0.5163105076990828]
# assert(np.any(np.isclose(e.data.item(), np.array(expected_energy))))
# assert(np.any(np.isclose(v.data.item(), np.array(expected_variance))))
def test3_wf_opt(self):
self.solver.sampler = self.sampler
self.solver.configure(track=['local_energy', 'parameters'],
loss='energy',
grad='auto')
obs = self.solver.run(5)
if __PLOT__:
plot_energy(obs.local_energy, e0=-1.1645, show_variance=True)
def test4_geo_opt(self):
self.solver.wf.ao.atom_coords[0, 2].data = torch.tensor(-0.37)
self.solver.wf.ao.atom_coords[1, 2].data = torch.tensor(0.37)
self.solver.configure(track=['local_energy'],
loss='energy',
grad='auto')
self.solver.geo_opt(5,
nepoch_wf_init=10,
nepoch_wf_update=5,
hdf5_group='geo_opt_correlated')
# load the best model
self.solver.wf.load(self.solver.hdf5file, 'geo_opt_correlated')
self.solver.wf.eval()
# sample and compute variables
obs = self.solver.single_point()
e, v = obs.energy, obs.variance
e = e.data.numpy()
v = v.data.numpy()
# it might be too much to assert with the ground state energy
assert (e > 2 * self.ground_state_energy and e < 0.)
assert (v > 0 and v < 2.)
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)
# wf
wfc = CorrelatedOrbital(mol,
kinetic='auto',
jastrow_type='pade_jastrow',
configs='single_double(2,4)',
include_all_mo=True)
# wf
wf = Orbital(mol,
kinetic='auto',
jastrow_type='pade_jastrow',
configs='single_double(2,4)',
include_all_mo=True)
random_fc_weight = torch.rand(wf.fc.weight.shape)
wf.fc.weight.data = random_fc_weight
wfc.fc.weight.data = random_fc_weight
nbatch = 3
class TestCorrelatedOrbitalWF(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 = CorrelatedOrbital(mol,
kinetic='auto',
configs='ground_state',
jastrow_type=FullyConnectedJastrow)
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, self.wf.nelec*3))
self.pos.requires_grad = True
def test_forward(self):
"""Value of the wave function."""
wfvals = self.wf(self.pos)
ref = torch.tensor([[0.2339], [0.2092], [0.3335], [0.2806], [0.1317],
[0.0996], [0.1210], [0.1406], [0.2626], [0.4675]])
# assert torch.allclose(wfvals.data, ref, rtol=1E-4, atol=1E-4)
def test_jacobian_mo(self):
"""Jacobian of the uncorrelated 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]
assert(torch.allclose(dmo.sum(-1),
dmo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_grad_mo(self):
"""Gradients of the uncorrelated MOs."""
mo = self.wf.pos2mo(self.pos)
dmo = self.wf.pos2mo(self.pos, derivative=1, jacobian=False)
dmo_grad = grad(
mo, self.pos, grad_outputs=torch.ones_like(mo))[0]
assert(torch.allclose(dmo.sum(-2),
dmo_grad.view(self.nbatch, self.wf.nelec, 3)))
def test_hess_mo(self):
"""Hessian of the uncorrelated MOs."""
mo = self.wf.pos2mo(self.pos)
d2mo = self.wf.pos2mo(self.pos, derivative=2)
d2mo_grad = hess(mo, self.pos)
assert(torch.allclose(d2mo.sum(-1),
d2mo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_jacobian_jast(self):
"""Jacobian of the jastrow values."""
jast = self.wf.ordered_jastrow(self.pos)
djast = self.wf.ordered_jastrow(self.pos, derivative=1)
djast_grad = grad(jast, self.pos,
grad_outputs=torch.ones_like(jast))[0]
assert(torch.allclose(djast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1),
djast.sum(-1)))
def test_grad_jast(self):
"""Gradients of the jastrow values."""
jast = self.wf.ordered_jastrow(self.pos)
djast = self.wf.ordered_jastrow(
self.pos, derivative=1, jacobian=False)
djast_grad = grad(jast, self.pos,
grad_outputs=torch.ones_like(jast))[0]
assert(torch.allclose(djast_grad.view(self.nbatch, self.wf.nelec, 3),
djast.sum(-2)))
def test_hess_jast(self):
"""Hessian of the jastrows."""
jast = self.wf.ordered_jastrow(self.pos)
d2jast = self.wf.ordered_jastrow(self.pos, derivative=2)
d2jast_grad = hess(jast, self.pos)
assert(torch.allclose(d2jast.sum(-1),
d2jast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_grad_cmo(self):
"""Gradients of the correlated MOs."""
cmo = self.wf.pos2cmo(self.pos)
dcmo = self.wf.get_gradient_operator(self.pos)
dcmo = dcmo.permute(1, 2, 3, 0)
shape = (self.nbatch, self.wf.nelec,
self.wf.nmo_opt, self.wf.nelec, 3)
dcmo = dcmo.reshape(*shape)
dcmo = dcmo.sum(2).sum(1)
dcmo_grad = grad(cmo, self.pos,
grad_outputs=torch.ones_like(cmo))[0]
dcmo_grad = dcmo_grad.reshape(self.nbatch, self.wf.nelec, 3)
assert(torch.allclose(dcmo, dcmo_grad))
def test_hess_cmo(self):
"""Hessian of the correlated MOs."""
val = self.wf.pos2cmo(self.pos)
d2val_grad = hess(val, self.pos)
d2val = self.wf.get_hessian_operator(self.pos)
d2val = d2val.permute(1, 2, 0, 3).sum(1)
assert(torch.allclose(d2val.sum(-1),
d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_jacobian_wf(self):
"""Jacobian of det(CMO).
\nabla det(CMOup) / det(CMOup) + \nabla det(CMOup) / det(CMOup) """
grad_jacobi = self.wf.gradients_jacobi(self.pos)
grad_auto = self.wf.gradients_autograd(self.pos)
assert(torch.allclose(grad_jacobi, grad_auto.sum(-1)))
def test_grad_wf(self):
"""Compute the gradients of the wf wrt to xyz coord of each elec."""
grad_jacobi = self.wf.gradients_jacobi(
self.pos, jacobian=False).squeeze()
grad_auto = self.wf.gradients_autograd(self.pos)
assert torch.allclose(grad_jacobi, grad_auto)
def test_kinetic_energy(self):
"""Kinetic energty."""
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_local_energy(self):
"""local energy."""
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)
class TestCorrelatedOrbitalWF(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 = CorrelatedOrbital(mol,
kinetic='auto',
jastrow_type='pade_jastrow',
configs='single_double(2,4)',
include_all_mo=True)
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.fc.weight.data = self.random_fc_weight
self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.weight.shape)
self.nbatch = 3
self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
def test_forward(self):
"""Value of the wave function."""
wfvals = self.wf(self.pos)
ref = torch.tensor([[-1.0935e-02], [6.4874e-02], [1.7879e-04],
[1.5797e-02], [7.4684e-02], [-4.4445e-02],
[-4.8149e-04], [-3.0355e-03], [-2.0027e-02],
[5.1957e-05]])
# assert torch.allclose(wfvals.data, ref, rtol=1E-4, atol=1E-4)
def test_jacobian_mo(self):
"""Jacobian of the uncorrelated 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]
assert (torch.allclose(
dmo.sum(-1),
dmo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_grad_mo(self):
"""Gradients of the uncorrelated MOs."""
mo = self.wf.pos2mo(self.pos)
dmo = self.wf.pos2mo(self.pos, derivative=1, jacobian=False)
dmo_grad = grad(mo, self.pos, grad_outputs=torch.ones_like(mo))[0]
assert (torch.allclose(dmo.sum(-2),
dmo_grad.view(self.nbatch, self.wf.nelec, 3)))
def test_hess_mo(self):
"""Hessian of the uncorrelated MOs."""
mo = self.wf.pos2mo(self.pos)
d2mo = self.wf.pos2mo(self.pos, derivative=2)
d2mo_grad = hess(mo, self.pos)
assert (torch.allclose(
d2mo.sum(-1),
d2mo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_jacobian_jast(self):
"""Jacobian of the jastrow values."""
jast = self.wf.ordered_jastrow(self.pos)
djast = self.wf.ordered_jastrow(self.pos, derivative=1)
djast_grad = grad(jast, self.pos,
grad_outputs=torch.ones_like(jast))[0]
assert (torch.allclose(
djast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1),
djast.sum(-1)))
def test_grad_jast(self):
"""Gradients of the jastrow values."""
jast = self.wf.ordered_jastrow(self.pos)
djast = self.wf.ordered_jastrow(self.pos, derivative=1, jacobian=False)
djast_grad = grad(jast, self.pos,
grad_outputs=torch.ones_like(jast))[0]
assert (torch.allclose(djast_grad.view(self.nbatch, self.wf.nelec, 3),
djast.sum(-2)))
def test_hess_jast(self):
"""Hessian of the jastrows."""
jast = self.wf.ordered_jastrow(self.pos)
d2jast = self.wf.ordered_jastrow(self.pos, derivative=2)
d2jast_grad = hess(jast, self.pos)
assert (torch.allclose(
d2jast.sum(-1),
d2jast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_grad_cmo(self):
"""Gradients of the correlated MOs."""
cmo = self.wf.pos2cmo(self.pos)
dcmo = self.wf.get_gradient_operator(self.pos)
dcmo = dcmo.permute(1, 2, 3, 0)
shape = (self.nbatch, self.wf.nelec, self.wf.nmo_opt, self.wf.nelec, 3)
dcmo = dcmo.reshape(*shape)
dcmo = dcmo.sum(2).sum(1)
dcmo_grad = grad(cmo, self.pos, grad_outputs=torch.ones_like(cmo))[0]
dcmo_grad = dcmo_grad.reshape(self.nbatch, self.wf.nelec, 3)
assert (torch.allclose(dcmo, dcmo_grad))
def test_hess_cmo(self):
"""Hessian of the correlated MOs."""
val = self.wf.pos2cmo(self.pos)
d2val_grad = hess(val, self.pos)
d2val = self.wf.get_hessian_operator(self.pos)
d2val = d2val.permute(1, 2, 0, 3).sum(1)
assert (torch.allclose(
d2val.sum(-1),
d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)))
def test_jacobian_wf(self):
"""Jacobian of det(CMO).
\nabla det(CMOup) / det(CMOup) + \nabla det(CMOup) / det(CMOup) """
grad_jacobi = self.wf.gradients_jacobi(self.pos)
grad_auto = self.wf.gradients_autograd(self.pos)
assert (torch.allclose(grad_jacobi, grad_auto.sum(-1)))
def test_grad_wf(self):
"""Compute the gradients of the wf wrt to xyz coord of each elec."""
grad_jacobi = self.wf.gradients_jacobi(self.pos,
jacobian=False).squeeze()
grad_auto = self.wf.gradients_autograd(self.pos)
assert torch.allclose(grad_jacobi, grad_auto)
def test_kinetic_energy(self):
"""Kinetic energty."""
eauto = self.wf.kinetic_energy_autograd(self.pos)
ejac = self.wf.kinetic_energy_jacobi(self.pos)
print(eauto)
print(ejac)
assert torch.allclose(eauto.data, ejac.data, rtol=1E-4, atol=1E-4)
def test_local_energy(self):
"""local energy."""
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)
hess[:, idim] = tmp[:, idim]
return hess
# molecule
mol = Molecule(
atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
wf = CorrelatedOrbital(
mol,
kinetic='jacobi',
jastrow_type='pade_jastrow',
configs='cas(2,2)')
wf.jastrow.weight.data = torch.rand(wf.jastrow.weight.shape)
nbatch = 10
pos = torch.tensor(np.random.rand(
nbatch, wf.nelec*3))
pos.requires_grad = True
cmo = wf.pos2cmo(pos)
sd = wf.pool(cmo)