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',
jastrow_kernel=PadeJastrowKernel,
include_all_mo=True,
configs='single_double(2,2)',
backflow_kernel=BackFlowKernelInverse,
orbital_dependent_backflow=False)
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.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 setUp(self):
torch.manual_seed(0)
np.random.seed(0)
set_torch_double_precision()
# molecule
self.mol = Molecule(atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# wave function
self.wf = SlaterJastrowBackFlow(self.mol,
kinetic='jacobi',
configs='single_double(2,2)',
backflow_kernel=BackFlowKernelPowerSum,
orbital_dependent_backflow=False,
include_all_mo=True)
# fc weights
self.wf.fc.weight.data = torch.rand(self.wf.fc.weight.shape)
# jastrow weights
self.wf.jastrow.jastrow_kernel.weight.data = torch.rand(
self.wf.jastrow.jastrow_kernel.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 = SolverSlaterJastrow(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# artificial pos
self.nbatch = 10
self.pos = torch.as_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()
# molecule
path_hdf5 = (
PATH_TEST / 'hdf5/LiH_adf_dz.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# wave function
self.wf = SlaterJastrowBackFlow(self.mol, kinetic='jacobi',
configs='single_double(2,2)',
orbital_dependent_backflow=True,
include_all_mo=True)
# fc weights
self.wf.fc.weight.data = torch.rand(self.wf.fc.weight.shape)
# jastrow weights
self.wf.jastrow.jastrow_kernel.weight.data = torch.rand(
self.wf.jastrow.jastrow_kernel.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 = SolverSlaterJastrow(wf=self.wf, sampler=self.sampler,
optimizer=self.opt)
# artificial pos
self.nbatch = 10
self.pos = torch.as_tensor(np.random.rand(
self.nbatch, self.wf.nelec*3))
self.pos.requires_grad = True
class MyBackflow(BackFlowKernelBase):
def __init__(self, mol, cuda, size=16):
super().__init__(mol, cuda)
self.fc1 = nn.Linear(1, size, bias=False)
self.fc2 = nn.Linear(size, 1, bias=False)
def forward(self, x):
original_shape = x.shape
x = x.reshape(-1, 1)
x = self.fc2(self.fc1(x))
return x.reshape(*original_shape)
# define the molecule
mol = Molecule(atom='Li 0. 0. 0.; H 3.14 0. 0.',
unit='angs',
calculator='pyscf',
basis='sto-3g',
name='LiH')
# define the wave function
wf = SlaterJastrowBackFlow(mol,
kinetic='jacobi',
backflow_kernel=MyBackflow,
backflow_kernel_kwargs={'size': 64},
configs='single_double(2,2)')
pos = torch.rand(10, wf.nelec * 3)
print(wf(pos))
class TestLiHBackFlowPySCF(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
set_torch_double_precision()
# molecule
self.mol = Molecule(atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# wave function
self.wf = SlaterJastrowBackFlow(self.mol,
kinetic='jacobi',
configs='single_double(2,2)',
orbital_dependent_backflow=True,
include_all_mo=True)
# fc weights
self.wf.fc.weight.data = torch.rand(self.wf.fc.weight.shape)
# jastrow weights
self.wf.jastrow.jastrow_kernel.weight.data = torch.rand(
self.wf.jastrow.jastrow_kernel.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 = SolverSlaterJastrow(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# artificial pos
self.nbatch = 10
self.pos = torch.as_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
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', 'parameters'],
loss='energy',
grad='manual')
obs = self.solver.run(5)
class TestSlaterJastrowBackFlow(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',
jastrow_kernel=PadeJastrowKernel,
include_all_mo=True,
configs='single_double(2,2)',
backflow_kernel=BackFlowKernelInverse,
orbital_dependent_backflow=False)
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.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):
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_jacobian_mo(self):
"""Jacobian of the BF 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(), dmo_grad.sum()))
psum_mo = dmo.sum(-1).sum(-1)
psum_mo_grad = dmo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)
psum_mo_grad = psum_mo_grad.T
assert (torch.allclose(psum_mo, psum_mo_grad))
def test_grad_mo(self):
"""Gradients of the BF MOs."""
mo = self.wf.pos2mo(self.pos)
dao = self.wf.ao(self.pos, derivative=1, sum_grad=False)
dmo = self.wf.ao2mo(dao)
dmo_grad = grad(mo, self.pos, grad_outputs=torch.ones_like(mo))[0]
assert (torch.allclose(dmo.sum(), dmo_grad.sum()))
dmo = dmo.sum(-1).sum(-1)
dmo_grad = dmo_grad.T
assert (torch.allclose(dmo, dmo_grad))
def test_hess_mo(self):
"""Hessian of the MOs."""
val = self.wf.pos2mo(self.pos)
d2val_grad = hess(val, self.pos)
d2ao = self.wf.ao(self.pos, derivative=2, sum_hess=False)
d2val = self.wf.ao2mo(d2ao)
assert (torch.allclose(d2val.sum(), d2val_grad.sum()))
d2val = d2val.reshape(4, 3, 5, 4, 6).sum(1).sum(-1).sum(-1)
d2val_grad = d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1)
d2val_grad = d2val_grad.T
assert (torch.allclose(d2val, d2val_grad))
def test_grad_wf(self):
pass
# grad_auto = self.wf.gradients_autograd(self.pos)
# grad_jac = self.wf.gradients_jacobi(self.pos)
# assert torch.allclose(
# grad_auto.data, grad_jac.data, rtol=1E-4, atol=1E-4)
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_jacobi
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)
print(ejac)
print(eauto)
assert torch.allclose(eauto.data, ejac.data, rtol=1E-4, atol=1E-4)
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)
def setUp(self):
set_torch_double_precision()
reset_generator()
# molecule
self.mol = Molecule(atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# molecule
self.mol_ref = Molecule(atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# backflow wave function
self.wf = SlaterJastrowBackFlow(self.mol,
kinetic='jacobi',
configs='single_double(2,2)',
include_all_mo=True)
self.wf.ao.backflow_trans.backflow_kernel.weight.data *= 0.
self.wf.ao.backflow_trans.backflow_kernel.weight.requires_grad = False
# normal wave function
self.wf_ref = SlaterJastrow(self.mol_ref,
kinetic='jacobi',
include_all_mo=True,
configs='single_double(2,2)')
# fc weights
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.fc.weight.data = self.random_fc_weight.clone()
self.wf_ref.fc.weight.data = self.random_fc_weight.clone()
# jastrow weights
self.random_jastrow_weight = torch.rand(
self.wf.jastrow.jastrow_kernel.weight.shape)
self.wf.jastrow.jastrow_kernel.weight.data = self.random_jastrow_weight.clone(
)
self.wf_ref.jastrow.jastrow_kernel.weight.data = self.random_jastrow_weight.clone(
)
reset_generator()
# sampler
self.sampler = Metropolis(nwalkers=5,
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'
})
reset_generator()
self.sampler_ref = Metropolis(nwalkers=5,
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
reset_generator()
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
reset_generator()
self.opt_ref = optim.Adam(self.wf_ref.parameters(), lr=0.01)
# solver
self.solver_ref = SolverSlaterJastrow(wf=self.wf_ref,
sampler=self.sampler_ref,
optimizer=self.opt_ref)
self.solver = SolverSlaterJastrow(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# artificial pos
self.nbatch = 10
self.pos = torch.as_tensor(
np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
class TestCompareLiHBackFlowPySCF(unittest.TestCase):
def setUp(self):
set_torch_double_precision()
reset_generator()
# molecule
self.mol = Molecule(atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# molecule
self.mol_ref = Molecule(atom='Li 0 0 0; H 0 0 3.015',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# backflow wave function
self.wf = SlaterJastrowBackFlow(self.mol,
kinetic='jacobi',
configs='single_double(2,2)',
include_all_mo=True)
self.wf.ao.backflow_trans.backflow_kernel.weight.data *= 0.
self.wf.ao.backflow_trans.backflow_kernel.weight.requires_grad = False
# normal wave function
self.wf_ref = SlaterJastrow(self.mol_ref,
kinetic='jacobi',
include_all_mo=True,
configs='single_double(2,2)')
# fc weights
self.random_fc_weight = torch.rand(self.wf.fc.weight.shape)
self.wf.fc.weight.data = self.random_fc_weight.clone()
self.wf_ref.fc.weight.data = self.random_fc_weight.clone()
# jastrow weights
self.random_jastrow_weight = torch.rand(
self.wf.jastrow.jastrow_kernel.weight.shape)
self.wf.jastrow.jastrow_kernel.weight.data = self.random_jastrow_weight.clone(
)
self.wf_ref.jastrow.jastrow_kernel.weight.data = self.random_jastrow_weight.clone(
)
reset_generator()
# sampler
self.sampler = Metropolis(nwalkers=5,
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'
})
reset_generator()
self.sampler_ref = Metropolis(nwalkers=5,
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
reset_generator()
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
reset_generator()
self.opt_ref = optim.Adam(self.wf_ref.parameters(), lr=0.01)
# solver
self.solver_ref = SolverSlaterJastrow(wf=self.wf_ref,
sampler=self.sampler_ref,
optimizer=self.opt_ref)
self.solver = SolverSlaterJastrow(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# artificial pos
self.nbatch = 10
self.pos = torch.as_tensor(
np.random.rand(self.nbatch, self.wf.nelec * 3))
self.pos.requires_grad = True
def test_0_wavefunction(self):
# compute the kinetic energy using bf orb
reset_generator()
e_bf = self.wf.kinetic_energy_jacobi(self.pos)
# compute the kinetic energy
reset_generator()
e_ref = self.wf_ref.kinetic_energy_jacobi(self.pos)
print(torch.stack([e_bf, e_ref], axis=1).squeeze())
assert torch.allclose(e_bf.data, e_ref.data, rtol=1E-4, atol=1E-4)
def test1_single_point(self):
# sample and compute observables
reset_generator()
obs = self.solver.single_point()
e_bf, v_bf = obs.energy, obs.variance
obs = self.solver.single_point()
e_bf, v_bf = obs.energy, obs.variance
# sample and compute observables
reset_generator()
obs_ref = self.solver_ref.single_point()
e_ref, v_ref = obs_ref.energy, obs.variance
obs_ref = self.solver_ref.single_point()
e_ref, v_ref = obs_ref.energy, obs.variance
# compare values
assert torch.allclose(e_bf.data, e_ref.data, rtol=1E-4, atol=1E-4)
assert torch.allclose(v_bf.data, v_ref.data, rtol=1E-4, atol=1E-4)
def test2_wf_opt_grad_auto(self):
nepoch = 5
# optimize using backflow
self.solver.configure(track=['local_energy'],
loss='energy',
grad='auto')
self.solver.configure_resampling(mode='never')
reset_generator()
obs = self.solver.run(nepoch)
e_bf = torch.as_tensor(np.array(obs.energy))
# optimize using ref
self.solver_ref.configure(track=['local_energy'],
loss='energy',
grad='auto')
self.solver_ref.configure_resampling(mode='never')
reset_generator()
obs_ref = self.solver_ref.run(nepoch)
e_ref = torch.as_tensor(np.array(obs_ref.energy))
assert torch.allclose(e_bf, e_ref, rtol=1E-4, atol=1E-4)
def test3_wf_opt_grad_manual(self):
nepoch = 5
# optimize using backflow
reset_generator()
self.solver.configure(track=['local_energy', 'parameters'],
loss='energy',
grad='manual')
obs = self.solver.run(nepoch)
e_bf = torch.as_tensor(np.array(obs.energy))
# optimize using backflow
reset_generator()
self.solver_ref.configure(track=['local_energy', 'parameters'],
loss='energy',
grad='manual')
obs = self.solver_ref.run(nepoch)
e_ref = torch.as_tensor(np.array(obs.energy))
# compare values
assert torch.allclose(e_bf, e_ref, rtol=1E-4, atol=1E-4)
