def setUp(self):
# define the molecule
at = 'C 0 0 0'
basis = 'dzp'
self.mol = Molecule(atom=at,
calculator='pyscf',
basis=basis,
unit='bohr')
self.m = gto.M(atom=at, basis=basis, unit='bohr')
# define the wave function
self.wf = Orbital(self.mol)
self.pos = torch.zeros(100, self.mol.nelec * 3)
self.pos[:, 0] = torch.linspace(-5, 5, 100)
self.pos[:, 1] = torch.linspace(-5, 5, 100)
self.pos[:, 2] = torch.linspace(-5, 5, 100)
self.pos = Variable(self.pos)
self.pos.requires_grad = True
self.x = self.pos[:, 0].detach().numpy()
class TestAOderivativesADF(unittest.TestCase):
def setUp(self):
# define the molecule
path_hdf5 = PATH_TEST / 'hdf5/C_adf_dzp.hdf5'
self.mol = Molecule(load=path_hdf5)
# define the wave function
self.wf = Orbital(self.mol, include_all_mo=True)
# define the grid points
npts = 11
self.pos = torch.rand(npts, self.mol.nelec * 3)
self.pos = Variable(self.pos)
self.pos.requires_grad = True
def test_ao_deriv(self):
ao = self.wf.ao(self.pos)
dao = self.wf.ao(self.pos, derivative=1)
dao_grad = grad(
ao, self.pos, grad_outputs=torch.ones_like(ao))[0]
gradcheck(self.wf.ao, self.pos)
assert(torch.allclose(dao.sum(), dao_grad.sum()))
def test_ao_hess(self):
ao = self.wf.ao(self.pos)
d2ao = self.wf.ao(self.pos, derivative=2)
d2ao_grad = hess(ao, self.pos)
assert(torch.allclose(d2ao.sum(), d2ao_grad.sum()))
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 = Orbital(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 = 1E-2 * 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)
# 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 setUp(self):
# define the molecule
path_hdf5 = PATH_TEST / 'hdf5/C_adf_dzp.hdf5'
self.mol = Molecule(load=path_hdf5)
# define the wave function
self.wf = Orbital(self.mol, include_all_mo=True)
# define the grid points
npts = 11
self.pos = torch.rand(npts, self.mol.nelec * 3)
self.pos = Variable(self.pos)
self.pos.requires_grad = True
class TestMOvaluesADF(unittest.TestCase):
def setUp(self):
# define the molecule
path_hdf5 = (PATH_TEST / 'hdf5/C_adf_dzp.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# define the wave function
self.wf = Orbital(self.mol, include_all_mo=True)
# define the grid points
self.npts = 21
pts = get_pts(self.npts)
self.pos = 10 * torch.ones(self.npts**2, self.mol.nelec * 3)
self.pos[:, :3] = pts
self.pos = Variable(self.pos)
self.pos.requires_grad = True
def test_mo(self):
movals = self.wf.mo_scf(self.wf.ao(self.pos)).detach().numpy()
for iorb in range(self.mol.basis.nmo):
path_cube = PATH_TEST / f'cube/C_MO_%SCF_A%{iorb + 1}.cub'
fname = path_cube.absolute().as_posix()
adf_ref_data = np.array(read_cubefile(fname)).reshape(
self.npts, self.npts)**2
qmctorch_data = (movals[:, 0, iorb]).reshape(self.npts,
self.npts)**2
delta = np.abs(adf_ref_data - qmctorch_data)
if __PLOT__:
plt.subplot(1, 3, 1)
plt.imshow(adf_ref_data)
plt.subplot(1, 3, 2)
plt.imshow(qmctorch_data)
plt.subplot(1, 3, 3)
plt.imshow(delta)
plt.show()
# the 0,0 point is much larger due to num instabilities
delta = np.sort(delta.flatten())
delta = delta[:-1]
assert (delta.mean() < 1E-3)
def setUp(self):
# molecule
self.mol = Molecule(
atom='H 0 0 -0.69; H 0 0 0.69',
unit='bohr',
calculator='pyscf',
basis='dzp')
# wave function
self.wf = Orbital(self.mol, kinetic='jacobi',
configs='single(2,2)',
use_jastrow=True)
npts = 51
self.pos = torch.zeros(npts, 6)
self.pos[:, 2] = torch.linspace(-2, 2, npts)
def setUp(self):
# define the molecule
path_hdf5 = (PATH_TEST / 'hdf5/C_adf_dzp.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# define the wave function
self.wf = Orbital(self.mol, include_all_mo=True)
# define the grid points
self.npts = 21
pts = get_pts(self.npts)
self.pos = torch.zeros(self.npts**2, self.mol.nelec * 3)
self.pos[:, :3] = pts
self.pos = Variable(self.pos)
self.pos.requires_grad = True
def setUp(self):
# define the molecule
at = 'C 0 0 0'
basis = 'dzp'
self.mol = Molecule(atom=at,
calculator='pyscf',
basis=basis,
unit='bohr')
self.m = gto.M(atom=at, basis=basis, unit='bohr')
# define the wave function
self.wf = Orbital(self.mol, include_all_mo=True)
# define the grid points
npts = 11
self.pos = torch.rand(npts, self.mol.nelec * 3)
self.pos = Variable(self.pos)
self.pos.requires_grad = True
class TestAOderivativesPyscf(unittest.TestCase):
def setUp(self):
# define the molecule
at = 'C 0 0 0'
basis = 'dzp'
self.mol = Molecule(atom=at,
calculator='pyscf',
basis=basis,
unit='bohr')
self.m = gto.M(atom=at, basis=basis, unit='bohr')
# define the wave function
self.wf = Orbital(self.mol, include_all_mo=True)
# define the grid points
npts = 11
self.pos = torch.rand(npts, self.mol.nelec * 3)
self.pos = Variable(self.pos)
self.pos.requires_grad = True
def test_ao_deriv(self):
ao = self.wf.ao(self.pos)
dao = self.wf.ao(self.pos, derivative=1)
dao_grad = grad(
ao, self.pos, grad_outputs=torch.ones_like(ao))[0]
gradcheck(self.wf.ao, self.pos)
assert(torch.allclose(dao.sum(), dao_grad.sum()))
def test_ao_hess(self):
ao = self.wf.ao(self.pos)
d2ao = self.wf.ao(self.pos, derivative=2)
d2ao_grad = hess(ao, self.pos)
assert(torch.allclose(d2ao.sum(), d2ao_grad.sum()))
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='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'
})
# 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)
path_hdf5 = (PATH_TEST / 'hdf5/CO2_adf_dzp.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# wave function
self.wf = Orbital(self.mol,
kinetic='jacobi',
configs='ground_state',
use_jastrow=True,
include_all_mo=False)
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 setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(atom='H 0 0 0; H 0 0 1.',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = Orbital(mol,
kinetic='auto',
include_all_mo=False,
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.pos = torch.tensor(np.random.rand(10, 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
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 setUp(self):
torch.manual_seed(0)
np.random.seed(0)
self.mol = Molecule(atom='C 0 0 0; O 0 0 2.190; O 0 0 -2.190',
calculator='pyscf',
basis='dzp',
unit='bohr')
# wave function
self.wf = Orbital(self.mol,
kinetic='jacobi',
configs='ground_state',
use_jastrow=True,
include_all_mo=False)
class TestAOvaluesADF(unittest.TestCase):
def setUp(self):
# define the molecule
path_hdf5 = (PATH_TEST / 'hdf5/C_adf_dzp.hdf5').absolute().as_posix()
self.mol = Molecule(load=path_hdf5)
# define the wave function
self.wf = Orbital(self.mol, include_all_mo=True)
# define the grid points
self.npts = 21
pts = get_pts(self.npts)
self.pos = torch.zeros(self.npts**2, self.mol.nelec * 3)
self.pos[:, :3] = pts
self.pos = Variable(self.pos)
self.pos.requires_grad = True
def test_ao(self):
aovals = self.wf.ao(self.pos).detach().numpy()
for iorb in range(self.mol.basis.nao):
path_cube = PATH_TEST / f'cube/C_AO_%Basis%AO{iorb}.cub'
fname = path_cube.absolute().as_posix()
adf_ref_data = np.array(read_cubefile(fname)).reshape(
self.npts, self.npts)
qmctorch_data = (aovals[:, 0, iorb]).reshape(self.npts, self.npts)
delta = np.abs(adf_ref_data - qmctorch_data)
if __PLOT__:
plt.subplot(1, 3, 1)
plt.imshow(adf_ref_data)
plt.subplot(1, 3, 2)
plt.imshow(qmctorch_data)
plt.subplot(1, 3, 3)
plt.imshow(delta)
plt.show()
assert (delta.mean() < 1E-3)
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(atom='C 0 0 0',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = Orbital(mol, kinetic='auto',
configs='ground_state').gto2sto()
self.pos = -0.25 + 0.5 * torch.tensor(np.random.rand(10, 18))
self.pos.requires_grad = True
class TestInterpolate(unittest.TestCase):
def setUp(self):
# molecule
self.mol = Molecule(
atom='H 0 0 -0.69; H 0 0 0.69',
unit='bohr',
calculator='pyscf',
basis='dzp')
# wave function
self.wf = Orbital(self.mol, kinetic='jacobi',
configs='single(2,2)',
use_jastrow=True)
npts = 51
self.pos = torch.zeros(npts, 6)
self.pos[:, 2] = torch.linspace(-2, 2, npts)
def test_ao(self):
interp_ao = InterpolateAtomicOrbitals(self.wf)
inter = interp_ao(self.pos)
ref = self.wf.ao(self.pos)
delta = (inter - ref).abs().mean()
assert(delta < 0.1)
def test_mo_reg(self):
interp_mo = InterpolateMolecularOrbitals(self.wf)
inter = interp_mo(self.pos, method='reg')
ref = self.wf.mo(self.wf.mo_scf(self.wf.ao(self.pos)))
delta = (inter - ref).abs().mean()
assert(delta < 0.1)
def test_mo_irreg(self):
interp_mo = InterpolateMolecularOrbitals(self.wf)
inter = interp_mo(self.pos, method='irreg')
ref = self.wf.mo(self.wf.mo_scf(self.wf.ao(self.pos)))
delta = (inter - ref).abs().mean()
assert(delta < 0.1)
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)
class TestOrbitalWF(unittest.TestCase):
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(atom='H 0 0 0; H 0 0 1.',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = Orbital(mol,
kinetic='auto',
include_all_mo=False,
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.pos = torch.tensor(np.random.rand(10, self.wf.nelec * 3))
self.pos.requires_grad = True
def test_forward(self):
wfvals = self.wf(self.pos)
ref = torch.tensor([[0.0977], [0.0618], [0.0587], [0.0861], [0.0466],
[0.0406], [0.0444], [0.0728], [0.0809], [0.1868]])
# assert torch.allclose(wfvals.data, ref, rtol=1E-4, atol=1E-4)
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)
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, kinpool=False)
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 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 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)
from qmctorch.scf import Molecule
from qmctorch.wavefunction import Orbital
from qmctorch.sampler import Metropolis
from qmctorch.solver import SolverOrbital
from qmctorch.utils import plot_walkers_traj
# define the molecule
mol = Molecule(atom='water.xyz',
unit='angs',
calculator='pyscf',
basis='sto-3g',
name='water')
# define the wave function
wf = Orbital(mol, kinetic='jacobi', configs='ground_State', use_jastrow=True)
# sampler
sampler = Metropolis(nwalkers=100,
nstep=500,
step_size=0.25,
nelec=wf.nelec,
ndim=wf.ndim,
init=mol.domain('atomic'),
move={
'type': 'one-elec',
'proba': 'normal'
})
# solver
solver = SolverOrbital(wf=wf, sampler=sampler)
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 = Orbital(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 = 1E-2 * 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_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, kinpool=False)
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)
# optimal H positions +0.69 and -0.69
# ground state energy : -31.688 eV -> -1.16 hartree
# bond dissociation energy 4.478 eV -> 0.16 hartree
set_torch_double_precision()
# define the molecule
mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
calculator='adf',
basis='dzp',
unit='bohr')
# define the wave function
wf = Orbital(mol,
kinetic='jacobi',
configs='cas(2,2)',
use_jastrow=True,
cuda=True)
wf.jastrow.weight.data[0] = 1.
# sampler
sampler = Metropolis(nwalkers=2000,
nstep=2000,
step_size=0.2,
ntherm=-1,
ndecor=100,
nelec=wf.nelec,
init=mol.domain('atomic'),
move={
'type': 'all-elec',
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)
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
pos = torch.tensor(np.random.rand(nbatch, wf.nelec * 3))
pos.requires_grad = True
print(wf(pos))
print(wfc(pos))
# bond distance : 0.74 A -> 1.38 a
# optimal H positions +0.69 and -0.69
# ground state energy : -31.688 eV -> -1.16 hartree
# bond dissociation energy 4.478 eV -> 0.16 hartree
set_torch_double_precision()
# define the molecule
mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
calculator='adf',
basis='dzp',
unit='bohr')
# define the wave function
wf = Orbital(mol,
kinetic='jacobi',
configs='single_double(2,2)',
use_jastrow=True)
wf.jastrow.weight.data[0] = 1.
# 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'
class TestAOvaluesPyscf(unittest.TestCase):
def setUp(self):
# define the molecule
at = 'C 0 0 0'
basis = 'dzp'
self.mol = Molecule(atom=at,
calculator='pyscf',
basis=basis,
unit='bohr')
self.m = gto.M(atom=at, basis=basis, unit='bohr')
# define the wave function
self.wf = Orbital(self.mol)
self.pos = torch.zeros(100, self.mol.nelec * 3)
self.pos[:, 0] = torch.linspace(-5, 5, 100)
self.pos[:, 1] = torch.linspace(-5, 5, 100)
self.pos[:, 2] = torch.linspace(-5, 5, 100)
self.pos = Variable(self.pos)
self.pos.requires_grad = True
self.x = self.pos[:, 0].detach().numpy()
def test_ao(self):
nzlm = np.linalg.norm(self.m.cart2sph_coeff(), axis=1)
aovals = self.wf.ao(self.pos).detach().numpy() / nzlm
aovals_ref = self.m.eval_ao('GTOval_cart',
self.pos.detach().numpy()[:, :3])
for iorb in range(self.mol.basis.nao):
if __PLOT__:
plt.plot(self.x, aovals[:, 0, iorb])
plt.plot(self.x, aovals_ref[:, iorb])
plt.show()
assert np.allclose(aovals[:, 0, iorb], aovals_ref[:, iorb])
def test_ao_deriv(self):
nzlm = np.linalg.norm(self.m.cart2sph_coeff(), axis=1)
daovals = self.wf.ao(self.pos, derivative=1).detach().numpy() / nzlm
daovals_ref = self.m.eval_gto('GTOval_ip_cart',
self.pos.detach().numpy()[:, :3])
daovals_ref = daovals_ref.sum(0)
for iorb in range(self.mol.basis.nao):
if __PLOT__:
plt.plot(self.x, daovals[:, 0, iorb])
plt.plot(self.x, daovals_ref[:, iorb])
plt.show()
assert np.allclose(daovals[:, 0, iorb], daovals_ref[:, iorb])