Creating your own Jastrow Factor

We present here how to create your own electron-electron Jastrow factor and use it in QMCTorch. During the import you must import the base class of the electron-electron Jastrow. We aso create a H2 molecule

from qmctorch.scf import Molecule
from qmctorch.wavefunction import SlaterJastrow
from qmctorch.wavefunction.jastrows.elec_elec.kernels import JastrowKernelElectronElectronBase
mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', unit='bohr', redo_scf=True)

INFO:QMCTorch|
INFO:QMCTorch| SCF Calculation
INFO:QMCTorch|  Removing H2_pyscf_sto-3g.hdf5 and redo SCF calculations
INFO:QMCTorch|  Running scf  calculation
converged SCF energy = -1.06599946214331
INFO:QMCTorch|  Molecule name       : H2
INFO:QMCTorch|  Number of electrons : 2
INFO:QMCTorch|  SCF calculator      : pyscf
INFO:QMCTorch|  Basis set           : sto-3g
INFO:QMCTorch|  SCF                 : HF
INFO:QMCTorch|  Number of AOs       : 2
INFO:QMCTorch|  Number of MOs       : 2
INFO:QMCTorch|  SCF Energy          : -1.066 Hartree

We can then use this base class to create a new Jastrow Factor. This is done in the same way one would create a new neural network layer in pytorch.

from torch import nn
class MyJastrow(JastrowKernelElectronElectronBase):
    def __init__(self, nup, ndown, cuda, size=16):
        super().__init__(nup, ndown, cuda)
        self.fc1 = nn.Linear(1, size, bias=False)
        self.fc2 = nn.Linear(size, 1, bias=False)
    def forward(self, x):
        nbatch, npair = x.shape
        x = x.reshape(-1,1)
        x = self.fc2(self.fc1(x))
        return x.reshape(nbatch, npair)

As seen above the prototype of the class constructor must be:

def __init__(self, nup, ndown, cuda, **kwargs)

The list of keyword argument can contain any pairs such as size=16.

This Jastrow use two fully connected layers. The size of the hidden layer is here controlled by a keyword argument size whose defauilt value is 16 It is important to note that the calculation of the first and second derivative of the jastrow kernel wrt the electronic positions are then done via automatic differentiation as implemented in the JastrowKernelElectronElectronBase class. Hence there is no need to derive and implement these derivatives. However it is necessary that the forward function, which takes as input a torch.tensor of dimension [Nbatch, Npair] first reshape this tensor to [Nbatch*Npair,1], then applies the transformation on this tensor and finally reshape the output tensor to [Nbatch, Npair].

To use this new Jastrow in the SlaterJastrow wave function ansatz we simply pass the class name as argument of the jastrow_kernel keyword argument. It is also possible to specify the values of the keyword argument size with the jastrow_kernel_kwargs. As seen below the pair of keyword argument and its value is passed as a python dictionary :

wf = SlaterJastrow(mol, jastrow_kernel=MyJastrow, jastrow_kernel_kwargs={'size' : 64})

INFO:QMCTorch|
INFO:QMCTorch| SCF Calculation
INFO:QMCTorch|  Running scf  calculation
converged SCF energy = -1.06599946214331
INFO:QMCTorch|  Molecule name       : H2
INFO:QMCTorch|  Number of electrons : 2
INFO:QMCTorch|  SCF calculator      : pyscf
INFO:QMCTorch|  Basis set           : sto-3g
INFO:QMCTorch|  SCF                 : HF
INFO:QMCTorch|  Number of AOs       : 2
INFO:QMCTorch|  Number of MOs       : 2
INFO:QMCTorch|  SCF Energy          : -1.066 Hartree
INFO:QMCTorch|
INFO:QMCTorch| Wave Function
INFO:QMCTorch|  Jastrow factor      : True
INFO:QMCTorch|  Jastrow kernel      : MyJastrow
INFO:QMCTorch|  Highest MO included : 2
INFO:QMCTorch|  Configurations      : ground_state
INFO:QMCTorch|  Number of confs     : 1
INFO:QMCTorch|  Kinetic energy      : jacobi
INFO:QMCTorch|  Number var  param   : 145
INFO:QMCTorch|  Cuda support        : False