def test3_domain(self):
mol = Molecule(load='H2_pyscf_sto-3g.hdf5')
domain_center = mol.domain('center')
assert (domain_center['center'] == np.array([0., 0., 0.5])).all()
domain_uniform = mol.domain('uniform')
assert domain_uniform == {'method': 'uniform', 'min': -0.5, 'max': 1.5}
domain_normal = mol.domain('normal')
assert np.all(domain_normal['mean'] == np.array([0., 0., 0.5]))
domain_atomic = mol.domain('atomic')
assert np.all(domain_atomic['atom_coords'] == np.array([[0., 0., 0.],
[0., 0., 1.]]))
class TestH2Stat(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
# 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 = SlaterJastrow(self.mol,
kinetic='jacobi',
configs='single(2,2)')
# sampler
self.sampler = Metropolis(nwalkers=100,
nstep=500,
step_size=0.5,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
ntherm=0,
ndecor=1,
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)
def test_sampling_traj(self):
pos = self.solver.sampler(self.solver.wf.pdf)
obs = self.solver.sampling_traj(pos)
plot_walkers_traj(obs.local_energy)
plot_block(obs.local_energy)
def test_stat(self):
pos = self.solver.sampler(self.solver.wf.pdf)
obs = self.solver.sampling_traj(pos)
if __PLOT__:
plot_blocking_energy(obs.local_energy, block_size=10)
plot_correlation_coefficient(obs.local_energy)
plot_integrated_autocorrelation_time(obs.local_energy)
class TestH2ADF(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
# molecule
path_hdf5 = (PATH_TEST / 'hdf5/H2_adf_dzp.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# wave function
self.wf = SlaterJastrow(self.mol,
kinetic='auto',
configs='single(2,2)')
# 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 = SolverSlaterJastrow(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_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
# vals on different archs
expected_energy = [-1.1572532653808594, -1.1501641653648578]
expected_variance = [0.05085879936814308, 0.05094174843043177]
assert (np.any(np.isclose(e.data.item(), np.array(expected_energy))))
assert (np.any(np.isclose(v.data.item(), np.array(expected_variance))))
# assert(e > 2 * self.ground_state_energy and e < 0.)
# assert(v > 0 and v < 5.)
def test_wf_opt_auto_grad(self):
self.solver.configure(track=['local_energy'],
loss='energy',
grad='auto')
obs = self.solver.run(5)
def test_wf_opt_manual_grad(self):
self.solver.configure(track=['local_energy'],
loss='energy',
grad='manual')
obs = self.solver.run(5)
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 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 = SlaterOrbitalDependentJastrow(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
for ker in self.wf.jastrow.jastrow_kernel.jastrow_functions:
ker.weight.data = torch.rand(1)
# 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)
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()
_, _ = 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'],
loss='energy',
grad='manual')
obs = self.solver.run(5)
class TestLiH(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
# 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 = Orbital(self.mol,
kinetic='jacobi',
configs='single(2,2)',
use_jastrow=True,
include_all_mo=False)
# 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)
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)
calculator='pyscf',
basis='sto-3g',
unit='bohr',
rank=hvd.local_rank())
# define the wave function
wf = SlaterJastrow(mol, kinetic='jacobi', configs='cas(2,2)', cuda=False)
# sampler
sampler = Metropolis(nwalkers=200,
nstep=200,
step_size=0.2,
ntherm=-1,
ndecor=100,
nelec=wf.nelec,
init=mol.domain('atomic'),
move={
'type': 'all-elec',
'proba': 'normal'
})
# optimizer
lr_dict = [{
'params': wf.jastrow.parameters(),
'lr': 3E-3
}, {
'params': wf.ao.parameters(),
'lr': 1E-6
}, {
'params': wf.mo.parameters(),
'lr': 1E-3
class TestH2Hvd(unittest.TestCase):
def setUp(self):
hvd.init()
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',
rank=hvd.local_rank())
# wave function
self.wf = SlaterJastrow(self.mol, kinetic='jacobi',
configs='cas(2,2)',
cuda=False)
# sampler
self.sampler = Metropolis(
nwalkers=200,
nstep=200,
step_size=0.2,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('atomic'),
move={
'type': 'all-elec',
'proba': 'normal'})
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# solver
self.solver = SolverSlaterJastrowHorovod(wf=self.wf, sampler=self.sampler,
optimizer=self.opt, rank=hvd.rank())
# ground state energy
self.ground_state_energy = -1.16
# ground state pos
self.ground_state_pos = 0.69
def test_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
# sample and compute observables
obs = self.solver.single_point()
e, v = obs.energy, obs.variance
e = e.data.item()
v = v.data.item()
assert np.isclose(e, -1.15, 0.2)
assert 0 < v < 2
def test_wf_opt(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.configure(track=['local_energy'], freeze=['ao', 'mo'],
loss='energy', grad='auto',
ortho_mo=False, clip_loss=False,
resampling={'mode': 'update',
'resample_every': 1,
'nstep_update': 50})
self.solver.run(10)
MPI.COMM_WORLD.barrier()
self.solver.wf.load(self.solver.hdf5file, 'wf_opt')
self.solver.wf.eval()
obs = self.solver.single_point()
e, v = obs.energy, obs.variance
e = e.data.numpy()
v = v.data.numpy()
assert np.isclose(e, -1.15, 0.2)
assert 0 < v < 2
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.)
class TestSampler(unittest.TestCase):
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
self.mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# orbital
self.wf = Orbital(self.mol)
def test_walkers_init(self):
"""Test different initialization methods of the walkers."""
w1 = Walkers(nwalkers=10,
nelec=self.mol.nelec,
ndim=3,
init=self.mol.domain('center'))
w2 = Walkers(nwalkers=10,
nelec=self.mol.nelec,
ndim=3,
init=self.mol.domain('uniform'))
w3 = Walkers(nwalkers=10,
nelec=self.mol.nelec,
ndim=3,
init=self.mol.domain('normal'))
w4 = Walkers(nwalkers=10,
nelec=self.mol.nelec,
ndim=3,
init=self.mol.domain('atomic'))
def test_metropolis(self):
"""Test Metropolis sampling."""
sampler = Metropolis(nwalkers=10,
nstep=20,
step_size=0.5,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'))
for m in ['one-elec', 'all-elec', 'all-elec-iter']:
for p in ['normal', 'uniform']:
sampler.configure_move({'type': m, 'proba': p})
pos = sampler(self.wf.pdf)
def test_hmc(self):
"""Test HMC sampler."""
sampler = Hamiltonian(nwalkers=10,
nstep=20,
step_size=0.1,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'))
pos = sampler(self.wf.pdf)
def test_gmh(self):
"""Test generalized MH."""
sampler = GeneralizedMetropolis(nwalkers=10,
nstep=20,
step_size=0.2,
nelec=self.wf.nelec,
ndim=self.wf.ndim,
init=self.mol.domain('normal'))
pos = sampler(self.wf.pdf)
class TestH2(unittest.TestCase):
def setUp(self):
torch.manual_seed(0)
np.random.seed(0)
# 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 = SlaterJastrow(self.mol,
kinetic='auto',
configs='single(2,2)')
# 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'
})
self.hmc_sampler = Hamiltonian(nwalkers=100,
nstep=200,
step_size=0.1,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('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)
# ground state energy
self.ground_state_energy = -1.16
# ground state pos
self.ground_state_pos = 0.69
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.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_single_point_hmc(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.hmc_sampler
# sample and compute observables
obs = self.solver.single_point()
e, v = obs.energy, obs.variance
# values on different arch
expected_energy = [-1.0877732038497925, -1.088576]
# values on different arch
expected_variance = [0.14341972768306732, 0.163771]
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.as_tensor(-0.37)
self.solver.wf.ao.atom_coords[1, 2].data = torch.as_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)
# load the best model
self.solver.wf.load(self.solver.hdf5file, 'geo_opt')
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.)
def test5_sampling_traj(self):
self.solver.sampler = self.sampler
self.solver.sampler.nstep = 100
self.solver.sampler.ntherm = 0
self.solver.sampler.ndecor = 1
pos = self.solver.sampler(self.solver.wf.pdf)
obs = self.solver.sampling_traj(pos)
if __PLOT__:
plot_walkers_traj(obs.local_energy)
plot_block(obs.local_energy)
plot_blocking_energy(obs.local_energy, block_size=10)
plot_correlation_coefficient(obs.local_energy)
plot_integrated_autocorrelation_time(obs.local_energy)
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)