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 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 TestLiH(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 = 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 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'], loss='energy', grad='manual') obs = self.solver.run(5)
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 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 = 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 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 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)