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 = SlaterJastrow(mol,
kinetic='auto',
include_all_mo=True,
configs='cas(2,2)')
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, mol.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 1.',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = SlaterJastrow(mol,
kinetic='auto',
configs='ground_state',
jastrow_kernel=FullyConnectedJastrowKernel)
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.as_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 = SlaterJastrowBackFlow(mol,
kinetic='jacobi',
include_all_mo=True,
configs='single_double(2,2)',
backflow_kernel=BackFlowKernelInverse,
orbital_dependent_backflow=False)
self.wf.ao.backflow_trans.backflow_kernel.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 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 = SlaterJastrow(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 TestSlaterJastrowElectronCusp(unittest.TestCase):
def setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(atom='He 0.5 0 0; He -0.5 0 0',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = SlaterJastrow(mol,
jastrow_kernel=FullyConnectedJastrowKernel,
kinetic='jacobi',
include_all_mo=True,
configs='ground_state')
self.nbatch = 100
def test_ee_cusp(self):
import matplotlib.pyplot as plt
pos_x = torch.Tensor(np.random.rand(self.nbatch, self.wf.nelec, 3))
x = torch.linspace(0, 2, self.nbatch)
pos_x[:, 0, :] = torch.as_tensor([0., 0., 0.]) + 1E-6
pos_x[:, 1, 0] = 0.
pos_x[:, 1, 1] = 0.
pos_x[:, 1, 2] = x
pos_x[:, 2, :] = 0.5 * torch.as_tensor([1., 1., 1.])
pos_x[:, 3, :] = -0.5 * torch.as_tensor([1., 1., 1.])
pos_x = pos_x.reshape(self.nbatch, self.wf.nelec * 3)
pos_x.requires_grad = True
x = x.detach().numpy()
j = self.wf.jastrow(pos_x).detach().numpy()
plt.plot(x, j)
plt.show()
dx = x[1] - x[0]
dj = (j[1:] - j[0:-1]) / dx
plt.plot(x[:-1], dj / j[:-1])
plt.show()
epot = self.wf.electronic_potential(pos_x).detach().numpy()
ekin = self.wf.kinetic_energy_jacobi(pos_x).detach().numpy()
eloc = self.wf.local_energy(pos_x).detach().numpy()
plt.plot(x, epot)
plt.plot(x, ekin)
plt.plot(x, eloc)
plt.show()
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 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 = SlaterJastrow(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 = SlaterJastrow(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):
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
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)
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 = SlaterJastrow(self.mol, kinetic='jacobi',
configs='single(2,2)')
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 = SlaterJastrow(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):
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 = SlaterJastrow(self.mol,
kinetic='jacobi',
jastrow_kernel=FullyConnectedJastrowKernel,
configs='single(2,2)',
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 = 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)
# 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 setUp(self):
torch.manual_seed(101)
np.random.seed(101)
set_torch_double_precision()
# molecule
mol = Molecule(atom='He 0.5 0 0; He -0.5 0 0',
unit='bohr',
calculator='pyscf',
basis='sto-3g',
redo_scf=True)
self.wf = SlaterJastrow(mol,
jastrow_kernel=FullyConnectedJastrowKernel,
kinetic='jacobi',
include_all_mo=True,
configs='ground_state')
self.nbatch = 100
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 = SlaterJastrow(self.mol, kinetic='jacobi',
configs='ground_state',
include_all_mo=False)
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 setUp(self):
torch.manual_seed(101)
np.random.seed(101)
# define the molecule
at = 'Li 0 0 0; H 0 0 1'
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 = SlaterJastrow(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 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 = SlaterJastrow(self.mol,
kinetic='jacobi',
configs='single(2,2)',
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 = SolverSlaterJastrow(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
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 = SlaterJastrow(self.mol,
kinetic='jacobi',
configs='ground_state',
include_all_mo=False)
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 = SlaterJastrow(self.mol)
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 = SlaterJastrow(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 = SlaterJastrow(mol, kinetic='auto',
configs='ground_state').gto2sto()
self.pos = -0.25 + 0.5 * \
torch.as_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 = SlaterJastrow(self.mol, kinetic='jacobi',
configs='single(2,2)')
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 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 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 = SlaterJastrow(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])
class TestOrbitalWF(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 = SlaterJastrow(mol,
kinetic='auto',
include_all_mo=True,
configs='cas(2,2)')
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, mol.nelec*3))
self.pos.requires_grad = True
def test_forward(self):
wfvals = self.wf(self.pos)
ref = torch.Tensor([[0.0522],
[0.0826],
[0.0774],
[0.1321],
[0.0459],
[0.0421],
[0.0551],
[0.0764],
[0.1164],
[0.2506]])
# assert torch.allclose(wfvals.data, ref, rtol=1E-4, atol=1E-4)
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))
# test spin down
pos_xdn = self.pos.clone()
perm_dn = list(range(self.wf.nelec))
perm_dn[self.wf.mol.nup-1] = self.wf.mol.nup
perm_dn[self.wf.mol.nup] = self.wf.mol.nup-1
pos_xdn = pos_xdn.reshape(self.nbatch, self.wf.nelec, 3)
pos_xdn = pos_xdn[:, perm_up, :].reshape(
self.nbatch, self.wf.nelec*3)
wfvals_xdn = self.wf(pos_xdn)
assert(torch.allclose(wfvals_ref, -1*wfvals_xdn))
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)
ref = torch.Tensor([[-1.6567],
[-0.8790],
[-2.8136],
[-0.3644],
[-0.4477],
[-0.2709],
[-0.6964],
[-0.3993],
[-0.4777],
[-0.0579]])
# assert torch.allclose(
# eloc_auto.data, ref, rtol=1E-4, atol=1E-4)
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)
ref = torch.Tensor([[0.6099],
[0.6438],
[0.6313],
[2.0512],
[0.0838],
[0.2699],
[0.5190],
[0.3381],
[1.8489],
[5.2226]])
# assert torch.allclose(
# ejac.data, ref, rtol=1E-4, atol=1E-4)
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 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 = SlaterJastrow(mol,
kinetic='auto',
configs='ground_state',
jastrow_kernel=FullyConnectedJastrowKernel)
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.as_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_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)
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)
grads = grads.reshape(10, self.wf.nelec, 3)
grad_auto = grad_auto.reshape(10, self.wf.nelec, 3)
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)
hvd.init()
if torch.cuda.is_available():
torch.cuda.set_device(hvd.rank())
set_torch_double_precision()
# define the molecule
mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
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
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
