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 TestCompareSlaterJastrowOrbitalDependentBackFlow(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 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=True) for ker in self.wf.ao.backflow_trans.backflow_kernel.orbital_dependent_kernel: ker.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 test_forward(self): """Check that backflow give same results as normal SlaterJastrow.""" wf_val = self.wf(self.pos) wf_val_ref = self.wf_ref(self.pos) assert (torch.allclose(wf_val, wf_val_ref)) def test_jacobian_mo(self): """Check that backflow give same results as normal SlaterJastrow.""" dmo = self.wf.pos2mo(self.pos, derivative=1) dmo_ref = self.wf_ref.pos2mo(self.pos, derivative=1) assert (torch.allclose(dmo.sum(0), dmo_ref)) def test_hess_mo(self): """Check that backflow give same results as normal SlaterJastrow.""" d2ao = self.wf.ao(self.pos, derivative=2, sum_hess=False) d2val = self.wf.ao2mo(d2ao) d2ao_ref = self.wf_ref.ao(self.pos, derivative=2, sum_hess=True) d2val_ref = self.wf_ref.ao2mo(d2ao_ref) assert (torch.allclose(d2val_ref, d2val.sum(0))) def test_grad_wf(self): pass def test_local_energy(self): self.wf.kinetic_energy = self.wf.kinetic_energy_jacobi eloc_jac = self.wf.local_energy(self.pos) self.wf_ref.kinetic_energy = self.wf_ref.kinetic_energy_jacobi eloc_jac_ref = self.wf_ref.local_energy(self.pos) assert torch.allclose(eloc_jac_ref.data, eloc_jac.data, rtol=1E-4, atol=1E-4) def test_kinetic_energy(self): ejac_ref = self.wf_ref.kinetic_energy_jacobi(self.pos) ejac = self.wf.kinetic_energy_jacobi(self.pos) assert torch.allclose(ejac_ref.data, ejac.data, rtol=1E-4, atol=1E-4)
class TestSlaterJastrowBackFlow(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 3.015', unit='bohr', calculator='pyscf', basis='sto-3g', redo_scf=True) self.wf = SlaterJastrowBackFlow(mol, kinetic='jacobi', jastrow_kernel=PadeJastrowKernel, include_all_mo=True, configs='single_double(2,2)', backflow_kernel=BackFlowKernelInverse, orbital_dependent_backflow=False) self.random_fc_weight = torch.rand(self.wf.fc.weight.shape) self.wf.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 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_jacobian_mo(self): """Jacobian of the BF 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] assert (torch.allclose(dmo.sum(), dmo_grad.sum())) psum_mo = dmo.sum(-1).sum(-1) psum_mo_grad = dmo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1) psum_mo_grad = psum_mo_grad.T assert (torch.allclose(psum_mo, psum_mo_grad)) def test_grad_mo(self): """Gradients of the BF MOs.""" mo = self.wf.pos2mo(self.pos) dao = self.wf.ao(self.pos, derivative=1, sum_grad=False) dmo = self.wf.ao2mo(dao) dmo_grad = grad(mo, self.pos, grad_outputs=torch.ones_like(mo))[0] assert (torch.allclose(dmo.sum(), dmo_grad.sum())) dmo = dmo.sum(-1).sum(-1) dmo_grad = dmo_grad.T assert (torch.allclose(dmo, dmo_grad)) def test_hess_mo(self): """Hessian of the MOs.""" val = self.wf.pos2mo(self.pos) d2val_grad = hess(val, self.pos) d2ao = self.wf.ao(self.pos, derivative=2, sum_hess=False) d2val = self.wf.ao2mo(d2ao) assert (torch.allclose(d2val.sum(), d2val_grad.sum())) d2val = d2val.reshape(4, 3, 5, 4, 6).sum(1).sum(-1).sum(-1) d2val_grad = d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1) d2val_grad = d2val_grad.T assert (torch.allclose(d2val, d2val_grad)) def test_grad_wf(self): pass # grad_auto = self.wf.gradients_autograd(self.pos) # grad_jac = self.wf.gradients_jacobi(self.pos) # assert torch.allclose( # grad_auto.data, grad_jac.data, rtol=1E-4, atol=1E-4) 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_jacobi 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) print(ejac) print(eauto) assert torch.allclose(eauto.data, ejac.data, rtol=1E-4, atol=1E-4)
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)