Exemplo n.º 1
0
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)
Exemplo n.º 2
0
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)
Exemplo n.º 3
0
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)
Exemplo n.º 4
0
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)