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 = CorrelatedOrbital(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 = torch.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 = CorrelatedOrbital(mol, kinetic='auto', jastrow_type='pade_jastrow', configs='single_double(2,4)', include_all_mo=True) self.random_fc_weight = torch.rand(self.wf.fc.weight.shape) self.wf.fc.weight.data = self.random_fc_weight self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.weight.shape) self.nbatch = 3 self.pos = 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) 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') # wave function self.wf = CorrelatedOrbital( self.mol, kinetic='auto', configs='cas(2,2)', jastrow_type=FullyConnectedJastrow, # jastrow_type='pade_jastrow', include_all_mo=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) set_torch_double_precision() # molecule path_hdf5 = (PATH_TEST / 'hdf5/LiH_adf_dz.hdf5').absolute().as_posix() self.mol = Molecule(load=path_hdf5) # wave function self.wf = CorrelatedOrbital(self.mol, kinetic='jacobi', configs='cas(2,2)', include_all_mo=True) # fc weights self.wf.fc.weight.data = torch.rand(self.wf.fc.weight.shape) # jastrow weights self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.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 = SolverOrbital(wf=self.wf, sampler=self.sampler, optimizer=self.opt) # artificial pos self.nbatch = 10 self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3)) self.pos.requires_grad = True
class TestLiHCorrelated(unittest.TestCase): def setUp(self): torch.manual_seed(0) np.random.seed(0) set_torch_double_precision() # molecule path_hdf5 = (PATH_TEST / 'hdf5/LiH_adf_dz.hdf5').absolute().as_posix() self.mol = Molecule(load=path_hdf5) # wave function self.wf = CorrelatedOrbital(self.mol, kinetic='jacobi', configs='cas(2,2)', include_all_mo=True) # fc weights self.wf.fc.weight.data = torch.rand(self.wf.fc.weight.shape) # jastrow weights self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.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 = SolverOrbital(wf=self.wf, sampler=self.sampler, optimizer=self.opt) # artificial pos self.nbatch = 10 self.pos = torch.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 # # 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 TestH2Correlated(unittest.TestCase): def setUp(self): 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') # wave function self.wf = CorrelatedOrbital( self.mol, kinetic='auto', configs='cas(2,2)', jastrow_type=FullyConnectedJastrow, # jastrow_type='pade_jastrow', include_all_mo=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_0_wavefunction(self): # artificial pos self.nbatch = 10 self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3)) self.pos.requires_grad = True 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): 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.1286007165908813, # -1.099538658544285] # # values on different arch # expected_variance = [0.45748308300971985, # 0.5163105076990828] # 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, hdf5_group='geo_opt_correlated') # load the best model self.solver.wf.load(self.solver.hdf5file, 'geo_opt_correlated') 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.)
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) # 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
class TestCorrelatedOrbitalWF(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 = CorrelatedOrbital(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 = torch.tensor(np.random.rand( self.nbatch, self.wf.nelec*3)) self.pos.requires_grad = True def test_forward(self): """Value of the wave function.""" wfvals = self.wf(self.pos) ref = torch.tensor([[0.2339], [0.2092], [0.3335], [0.2806], [0.1317], [0.0996], [0.1210], [0.1406], [0.2626], [0.4675]]) # assert torch.allclose(wfvals.data, ref, rtol=1E-4, atol=1E-4) def test_jacobian_mo(self): """Jacobian of the uncorrelated 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(-1), dmo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_grad_mo(self): """Gradients of the uncorrelated MOs.""" mo = self.wf.pos2mo(self.pos) dmo = self.wf.pos2mo(self.pos, derivative=1, jacobian=False) dmo_grad = grad( mo, self.pos, grad_outputs=torch.ones_like(mo))[0] assert(torch.allclose(dmo.sum(-2), dmo_grad.view(self.nbatch, self.wf.nelec, 3))) def test_hess_mo(self): """Hessian of the uncorrelated MOs.""" mo = self.wf.pos2mo(self.pos) d2mo = self.wf.pos2mo(self.pos, derivative=2) d2mo_grad = hess(mo, self.pos) assert(torch.allclose(d2mo.sum(-1), d2mo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_jacobian_jast(self): """Jacobian of the jastrow values.""" jast = self.wf.ordered_jastrow(self.pos) djast = self.wf.ordered_jastrow(self.pos, derivative=1) djast_grad = grad(jast, self.pos, grad_outputs=torch.ones_like(jast))[0] assert(torch.allclose(djast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1), djast.sum(-1))) def test_grad_jast(self): """Gradients of the jastrow values.""" jast = self.wf.ordered_jastrow(self.pos) djast = self.wf.ordered_jastrow( self.pos, derivative=1, jacobian=False) djast_grad = grad(jast, self.pos, grad_outputs=torch.ones_like(jast))[0] assert(torch.allclose(djast_grad.view(self.nbatch, self.wf.nelec, 3), djast.sum(-2))) def test_hess_jast(self): """Hessian of the jastrows.""" jast = self.wf.ordered_jastrow(self.pos) d2jast = self.wf.ordered_jastrow(self.pos, derivative=2) d2jast_grad = hess(jast, self.pos) assert(torch.allclose(d2jast.sum(-1), d2jast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_grad_cmo(self): """Gradients of the correlated MOs.""" cmo = self.wf.pos2cmo(self.pos) dcmo = self.wf.get_gradient_operator(self.pos) dcmo = dcmo.permute(1, 2, 3, 0) shape = (self.nbatch, self.wf.nelec, self.wf.nmo_opt, self.wf.nelec, 3) dcmo = dcmo.reshape(*shape) dcmo = dcmo.sum(2).sum(1) dcmo_grad = grad(cmo, self.pos, grad_outputs=torch.ones_like(cmo))[0] dcmo_grad = dcmo_grad.reshape(self.nbatch, self.wf.nelec, 3) assert(torch.allclose(dcmo, dcmo_grad)) def test_hess_cmo(self): """Hessian of the correlated MOs.""" val = self.wf.pos2cmo(self.pos) d2val_grad = hess(val, self.pos) d2val = self.wf.get_hessian_operator(self.pos) d2val = d2val.permute(1, 2, 0, 3).sum(1) assert(torch.allclose(d2val.sum(-1), d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_jacobian_wf(self): """Jacobian of det(CMO). \nabla det(CMOup) / det(CMOup) + \nabla det(CMOup) / det(CMOup) """ grad_jacobi = self.wf.gradients_jacobi(self.pos) grad_auto = self.wf.gradients_autograd(self.pos) assert(torch.allclose(grad_jacobi, grad_auto.sum(-1))) def test_grad_wf(self): """Compute the gradients of the wf wrt to xyz coord of each elec.""" grad_jacobi = self.wf.gradients_jacobi( self.pos, jacobian=False).squeeze() grad_auto = self.wf.gradients_autograd(self.pos) assert torch.allclose(grad_jacobi, grad_auto) def test_kinetic_energy(self): """Kinetic energty.""" 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_local_energy(self): """local energy.""" 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)
class TestCorrelatedOrbitalWF(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 = CorrelatedOrbital(mol, kinetic='auto', jastrow_type='pade_jastrow', configs='single_double(2,4)', include_all_mo=True) self.random_fc_weight = torch.rand(self.wf.fc.weight.shape) self.wf.fc.weight.data = self.random_fc_weight self.wf.jastrow.weight.data = torch.rand(self.wf.jastrow.weight.shape) self.nbatch = 3 self.pos = torch.tensor(np.random.rand(self.nbatch, self.wf.nelec * 3)) self.pos.requires_grad = True def test_forward(self): """Value of the wave function.""" wfvals = self.wf(self.pos) ref = torch.tensor([[-1.0935e-02], [6.4874e-02], [1.7879e-04], [1.5797e-02], [7.4684e-02], [-4.4445e-02], [-4.8149e-04], [-3.0355e-03], [-2.0027e-02], [5.1957e-05]]) # assert torch.allclose(wfvals.data, ref, rtol=1E-4, atol=1E-4) def test_jacobian_mo(self): """Jacobian of the uncorrelated 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(-1), dmo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_grad_mo(self): """Gradients of the uncorrelated MOs.""" mo = self.wf.pos2mo(self.pos) dmo = self.wf.pos2mo(self.pos, derivative=1, jacobian=False) dmo_grad = grad(mo, self.pos, grad_outputs=torch.ones_like(mo))[0] assert (torch.allclose(dmo.sum(-2), dmo_grad.view(self.nbatch, self.wf.nelec, 3))) def test_hess_mo(self): """Hessian of the uncorrelated MOs.""" mo = self.wf.pos2mo(self.pos) d2mo = self.wf.pos2mo(self.pos, derivative=2) d2mo_grad = hess(mo, self.pos) assert (torch.allclose( d2mo.sum(-1), d2mo_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_jacobian_jast(self): """Jacobian of the jastrow values.""" jast = self.wf.ordered_jastrow(self.pos) djast = self.wf.ordered_jastrow(self.pos, derivative=1) djast_grad = grad(jast, self.pos, grad_outputs=torch.ones_like(jast))[0] assert (torch.allclose( djast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1), djast.sum(-1))) def test_grad_jast(self): """Gradients of the jastrow values.""" jast = self.wf.ordered_jastrow(self.pos) djast = self.wf.ordered_jastrow(self.pos, derivative=1, jacobian=False) djast_grad = grad(jast, self.pos, grad_outputs=torch.ones_like(jast))[0] assert (torch.allclose(djast_grad.view(self.nbatch, self.wf.nelec, 3), djast.sum(-2))) def test_hess_jast(self): """Hessian of the jastrows.""" jast = self.wf.ordered_jastrow(self.pos) d2jast = self.wf.ordered_jastrow(self.pos, derivative=2) d2jast_grad = hess(jast, self.pos) assert (torch.allclose( d2jast.sum(-1), d2jast_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_grad_cmo(self): """Gradients of the correlated MOs.""" cmo = self.wf.pos2cmo(self.pos) dcmo = self.wf.get_gradient_operator(self.pos) dcmo = dcmo.permute(1, 2, 3, 0) shape = (self.nbatch, self.wf.nelec, self.wf.nmo_opt, self.wf.nelec, 3) dcmo = dcmo.reshape(*shape) dcmo = dcmo.sum(2).sum(1) dcmo_grad = grad(cmo, self.pos, grad_outputs=torch.ones_like(cmo))[0] dcmo_grad = dcmo_grad.reshape(self.nbatch, self.wf.nelec, 3) assert (torch.allclose(dcmo, dcmo_grad)) def test_hess_cmo(self): """Hessian of the correlated MOs.""" val = self.wf.pos2cmo(self.pos) d2val_grad = hess(val, self.pos) d2val = self.wf.get_hessian_operator(self.pos) d2val = d2val.permute(1, 2, 0, 3).sum(1) assert (torch.allclose( d2val.sum(-1), d2val_grad.view(self.nbatch, self.wf.nelec, 3).sum(-1))) def test_jacobian_wf(self): """Jacobian of det(CMO). \nabla det(CMOup) / det(CMOup) + \nabla det(CMOup) / det(CMOup) """ grad_jacobi = self.wf.gradients_jacobi(self.pos) grad_auto = self.wf.gradients_autograd(self.pos) assert (torch.allclose(grad_jacobi, grad_auto.sum(-1))) def test_grad_wf(self): """Compute the gradients of the wf wrt to xyz coord of each elec.""" grad_jacobi = self.wf.gradients_jacobi(self.pos, jacobian=False).squeeze() grad_auto = self.wf.gradients_autograd(self.pos) assert torch.allclose(grad_jacobi, grad_auto) def test_kinetic_energy(self): """Kinetic energty.""" eauto = self.wf.kinetic_energy_autograd(self.pos) ejac = self.wf.kinetic_energy_jacobi(self.pos) print(eauto) print(ejac) assert torch.allclose(eauto.data, ejac.data, rtol=1E-4, atol=1E-4) def test_local_energy(self): """local energy.""" 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)
hess[:, idim] = tmp[:, idim] return hess # molecule mol = Molecule( atom='Li 0 0 0; H 0 0 3.015', unit='bohr', calculator='pyscf', basis='sto-3g', redo_scf=True) wf = CorrelatedOrbital( mol, kinetic='jacobi', jastrow_type='pade_jastrow', configs='cas(2,2)') wf.jastrow.weight.data = torch.rand(wf.jastrow.weight.shape) nbatch = 10 pos = torch.tensor(np.random.rand( nbatch, wf.nelec*3)) pos.requires_grad = True cmo = wf.pos2cmo(pos) sd = wf.pool(cmo)