Merge pull request #134 from NLESC-JCER/unify_slaterjastrow

Unify slaterjastrow

.github/workflows/build.yml

@@ -26,17 +26,17 @@ jobs:
         conda-channels: anaconda
     - run: conda --version
     - run: which python
-    - run: conda install mpi4py h5py==3.1.0 pytorch torchvision cpuonly -c pytorch -c conda-forge
+    - run: conda install rdkit mpi4py h5py==3.1.0 pytorch torchvision cpuonly -c pytorch -c conda-forge
 
     - name: Install the package
       run: pip install .[test,hpc]
       env:
         CONDA_PREFIX: /usr/share/miniconda
 
-    - name: Test with multithreading
-      env:
-        CONDA_PREFIX: /usr/share/miniconda
-      run: mpirun -np 2 coverage run -m pytest tests_hvd
+    # - name: Test with multithreading
+    #   env:
+    #     CONDA_PREFIX: /usr/share/miniconda
+    #   run: mpirun -np 2 coverage run -m pytest tests_hvd
 
     - name: Test with single thread
       env:

bin/qmctorch_energy_plotter

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+import h5py
+import matplotlib.pyplot as plt
+import argparse
+
+
+def get_energy(mol):
+    with h5py.File(mol, 'r') as f5:
+        data = f5['wf_opt']['energy'][()]
+    return data
+
+
+def get_correlation_energy(e, e0, ehf):
+    return 1 - (e-e0)/(ehf-e0)
+
+
+def plot_percent_correlation_energy(args):
+
+    nepoch = args.num_epoch
+    energy = []
+    percent_correlation_energy = []
+    for mol in args.filename:
+        e = get_energy(mol)[:nepoch]
+        print(e[-1])
+        energy.append(e)
+        percent_correlation_energy.append(
+            get_correlation_energy(e, args.exact_energy, args.hf_energy))
+
+    plt_fn = plt.plot
+    if args.semi_logy:
+        plt_fn = plt.semilogy
+
+    for ec in percent_correlation_energy:
+        plt_fn(ec)
+    plt.show()
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('filename', nargs='+',
+                        help='name of the files')
+    parser.add_argument('-l', '--labels', nargs='+',
+                        help='label of the data')
+    parser.add_argument('-ne', '--num_epoch', type=int,
+                        default=-1, help='Number of epcoh to plot')
+    parser.add_argument('-e0', '--exact_energy', type=float,
+                        default=None, help='True exact energy of thre system')
+    parser.add_argument('-ehf', '--hf_energy', type=float,
+                        default=None, help='Hartree Fock energy of thre system')
+    parser.add_argument('-log', '--semi_logy', action='store_true',
+                        help='plot on semilog y axis')
+    args = parser.parse_args()
+
+    plot_percent_correlation_energy(args)

example/autocorrelation/h2.py

@@ -5,8 +5,8 @@
 from qmctorch.scf import Molecule
 from qmctorch.solver import SolverSlaterJastrow
 from qmctorch.utils import plot_correlation_coefficient, plot_integrated_autocorrelation_time
-from qmctorch.wavefunction import SlaterJastrow
-
+from qmctorch.wavefunction.slater_jastrow import SlaterJastrow
+from qmctorch.wavefunction.jastrows.elec_elec import JastrowFactor, PadeJastrowKernel
 torch.manual_seed(0)
 
 # molecule
@@ -16,8 +16,13 @@
     calculator='pyscf',
     basis='sto-3g')
 
-# wave function
+
+# jastrow
+jastrow = JastrowFactor(mol, PadeJastrowKernel)
+
+# wave funtion
 wf = SlaterJastrow(mol, kinetic='auto',
+                   jastrow=jastrow,
                    configs='single(2,2)')
 
 # sampler
@@ -43,5 +48,6 @@
 
 rho, tau = plot_correlation_coefficient(obs.local_energy)
 print(f'fit exp(-x/tau), tau={tau}')
-iat = plot_integrated_autocorrelation_time(obs.local_energy, rho=rho, C=5)
+iat = plot_integrated_autocorrelation_time(
+    obs.local_energy, rho=rho, C=5)
 print(f"integrated autocorrelation time: {iat}")

example/backflow/backflow.py

@@ -3,8 +3,13 @@
 from torch import nn
 
 from qmctorch.scf import Molecule
+
+from qmctorch.wavefunction.orbitals.backflow import BackFlowTransformation
 from qmctorch.wavefunction.orbitals.backflow.kernels import BackFlowKernelBase
-from qmctorch.wavefunction import SlaterJastrowBackFlow
+
+
+from qmctorch.wavefunction.slater_jastrow import SlaterJastrow
+from qmctorch.wavefunction.jastrows.elec_elec import JastrowFactor, PadeJastrowKernel
 
 
 class MyBackflow(BackFlowKernelBase):
@@ -25,11 +30,17 @@ def forward(self, x):
 mol = Molecule(atom='Li 0. 0. 0.; H 3.14 0. 0.', unit='angs',
                calculator='pyscf', basis='sto-3g', name='LiH')
 
+# jastrow
+jastrow = JastrowFactor(mol, PadeJastrowKernel)
+
+# backflow
+backflow = BackFlowTransformation(mol, MyBackflow, {'size': 64})
+
 # define the wave function
-wf = SlaterJastrowBackFlow(mol, kinetic='jacobi',
-                           backflow_kernel=MyBackflow,
-                           backflow_kernel_kwargs={'size': 64},
-                           configs='single_double(2,2)')
+wf = SlaterJastrow(mol, kinetic='jacobi',
+                   jastrow=jastrow,
+                   backflow=backflow,
+                   configs='single_double(2,2)')
 
 pos = torch.rand(10, wf.nelec*3)
 print(wf(pos))

example/gpu/h2.py

@@ -1,7 +1,8 @@
 from torch import optim
 
 from qmctorch.scf import Molecule
-from qmctorch.wavefunction import SlaterJastrow
+from qmctorch.wavefunction.slater_jastrow import SlaterJastrow
+from qmctorch.wavefunction.jastrows.elec_elec import JastrowFactor, PadeJastrowKernel
 from qmctorch.solver import SolverSlaterJastrow
 from qmctorch.sampler import Metropolis
 from qmctorch.utils import set_torch_double_precision
@@ -20,9 +21,14 @@
                basis='dzp',
                unit='bohr')
 
+
+# jastrow
+jastrow = JastrowFactor(mol, PadeJastrowKernel)
+
 # define the wave function
 wf = SlaterJastrow(mol, kinetic='jacobi',
                    configs='cas(2,2)',
+                   jastrow=jastrow,
                    cuda=True)
 
 # sampler

example/jast_graph.py

@@ -0,0 +1,19 @@
+
+from qmctorch.wavefunction.jastrows.graph.jastrow_graph import JastrowFactorGraph
+import torch
+from torch.autograd import grad
+nup = 2
+ndown = 2
+atomic_pos = torch.rand(2, 3)
+atom_types = ["Li", "H"]
+jast = JastrowFactorGraph(nup, ndown,
+                          atomic_pos,
+                          atom_types)
+
+
+pos = torch.rand(10, 12)
+pos.requires_grad = True
+jval = jast(pos)
+
+gval = jast(pos, derivative=1)
+hval = jast(pos, derivative=2)

example/optimization/h2.py

@@ -1,39 +1,47 @@
-
+import torch
+import numpy as np
 from torch import optim
 
 from qmctorch.scf import Molecule
-from qmctorch.wavefunction import SlaterJastrow
+
 from qmctorch.solver import SolverSlaterJastrow
 from qmctorch.sampler import Metropolis, Hamiltonian
 from qmctorch.utils import set_torch_double_precision
-from qmctorch.utils import (plot_energy, plot_data)
-
-from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel
+from qmctorch.utils.plot_data import (plot_energy, plot_data)
+from qmctorch.wavefunction.slater_jastrow import SlaterJastrow
+from qmctorch.wavefunction.jastrows.elec_elec import JastrowFactor, PadeJastrowKernel
 
 # bond distance : 0.74 A -> 1.38 a
 # optimal H positions +0.69 and -0.69
 # ground state energy : -31.688 eV -> -1.16 hartree
 # bond dissociation energy 4.478 eV -> 0.16 hartree
 
 set_torch_double_precision()
+torch.random.manual_seed(0)
+np.random.seed(0)
 
 # define the molecule
 mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
                calculator='pyscf',
                basis='sto-3g',
                unit='bohr')
 
+# jastrow
+jastrow = JastrowFactor(mol, PadeJastrowKernel)
+
 # define the wave function
 wf = SlaterJastrow(mol, kinetic='jacobi',
                    configs='single_double(2,2)',
-                   jastrow_kernel=PadeJastrowKernel)
+                   jastrow=jastrow)
 
 # sampler
-sampler = Hamiltonian(nwalkers=100, nstep=100, nelec=wf.nelec,
-                      step_size=0.1, L=30,
-                      ntherm=-1, ndecor=10,
-                      init=mol.domain('atomic'))
+# sampler = Hamiltonian(nwalkers=100, nstep=100, nelec=wf.nelec,
+#                       step_size=0.1, L=30,
+#                       ntherm=-1, ndecor=10,
+#                       init=mol.domain('atomic'))
 
+sampler = Metropolis(nwalkers=10, nstep=200, nelec=wf.nelec, ntherm=100, ndecor=10,
+                     step_size=0.05, init=mol.domain('atomic'))
 
 # optimizer
 lr_dict = [{'params': wf.jastrow.parameters(), 'lr': 3E-3},
@@ -50,19 +58,21 @@
                              optimizer=opt, scheduler=None)
 
 # perform a single point calculation
-obs = solver.single_point()
+# obs = solver.single_point()
 
 # configure the solver
 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})
+                             'nstep_update': 150,
+                             'ntherm_update': 50}
+                 )
 
 # optimize the wave function
-obs = solver.run(250)
+obs = solver.run(5)  # , batchsize=10)
 
 # plot
-plot_energy(obs.local_energy, e0=-1.1645, show_variance=True)
-plot_data(solver.observable, obsname='jastrow.weight')
+# plot_energy(obs.local_energy, e0=-1.1645, show_variance=True)
+# plot_data(solver.observable, obsname='jastrow.weight')

example/single_point/h2.py

@@ -1,5 +1,6 @@
 from qmctorch.scf import Molecule
-from qmctorch.wavefunction import SlaterJastrow
+from qmctorch.wavefunction.slater_jastrow import SlaterJastrow
+from qmctorch.wavefunction.jastrows.elec_elec import JastrowFactor, PadeJastrowKernel
 from qmctorch.sampler import Metropolis
 from qmctorch.solver import SolverSlaterJastrow
 from qmctorch.utils import plot_walkers_traj
@@ -10,10 +11,12 @@
 mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
                calculator='pyscf', basis='dzp', unit='bohr')
 
+# jastrow
+jastrow = JastrowFactor(mol, PadeJastrowKernel)
 
 # define the wave function
 wf = SlaterJastrow(mol, kinetic='jacobi',
-                   configs='ground_state').gto2sto()
+                   configs='ground_state', jastrow=jastrow).gto2sto()
 
 # sampler
 sampler = Metropolis(nwalkers=1000, nstep=1000, step_size=0.25,

example/single_point/h2o_sampling.py

@@ -1,5 +1,6 @@
 from qmctorch.scf import Molecule
-from qmctorch.wavefunction import SlaterJastrow
+from qmctorch.wavefunction.slater_jastrow import SlaterJastrow
+from qmctorch.wavefunction.jastrows.elec_elec import JastrowFactor, PadeJastrowKernel
 from qmctorch.sampler import Metropolis
 from qmctorch.solver import SolverSlaterJastrow
 from qmctorch.utils import plot_walkers_traj
@@ -9,9 +10,12 @@
 mol = Molecule(atom='water.xyz', unit='angs',
                calculator='pyscf', basis='sto-3g', name='water')
 
+# jastrow
+jastrow = JastrowFactor(mol, PadeJastrowKernel)
+
 # define the wave function
 wf = SlaterJastrow(mol, kinetic='jacobi',
-                   configs='ground_state')
+                   configs='ground_state', jastrow=jastrow)
 
 # sampler
 sampler = Metropolis(nwalkers=100, nstep=500, step_size=0.25,

h5x/baseimport.py

@@ -1,4 +1,4 @@
-from qmctorch.utils import (
+from qmctorch.utils.plot_data import (
     plot_energy, plot_data, plot_block, plot_walkers_traj)
 import matplotlib.pyplot as plt
 import numpy as np