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)
class TestH2ADFJacobi(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 = Orbital(self.mol, kinetic='jacobi',
configs='single(2,2)',
use_jastrow=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_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
print(e.data.item(), v.data.item())
# vals on different archs
expected_energy = [-1.1571345329284668,
-1.1501641653648578]
expected_variance = [0.05087674409151077,
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))))
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 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 = Orbital(self.mol, kinetic='jacobi',
configs='single(2,2)',
use_jastrow=True)
# 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 = SolverOrbital(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 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 = Orbital(self.mol,
kinetic='auto',
configs='single(2,2)',
use_jastrow=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'
})
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 = 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 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.077970027923584, -1.027975961270174]
# values on different arch
expected_variance = [0.17763596773147583, 0.19953053065068135]
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)
# 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)
