Combining several Jastrow Factors

It is often useful to use mutliple jastrow factors to go beyond the simple electron-electron Jastrow. We show here how to do that easily through our

import torch
from qmctorch.scf import Molecule
from qmctorch.wavefunction import SlaterJastrow

from qmctorch.wavefunction.jastrows.elec_elec import (
    JastrowFactor as JastrowFactorElecElec,
    FullyConnectedJastrowKernel as FCEE,
)
from qmctorch.wavefunction.jastrows.elec_nuclei import (
    JastrowFactor as JastrowFactorElecNuclei,
    FullyConnectedJastrowKernel as FCEN,
)

INFO:QMCTorch|  ____    __  ______________             _
INFO:QMCTorch| / __ \  /  |/  / ___/_  __/__  ________/ /
INFO:QMCTorch|/ /_/ / / /|_/ / /__  / / / _ \/ __/ __/ _ \
INFO:QMCTorch|\___\_\/_/  /_/\___/ /_/  \___/_/  \__/_//_/

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.

We first need a molecule

mol = Molecule(
        atom="Li 0 0 0; H 0 0 3.14",
        unit='bohr',
        calculator="pyscf",
        basis="sto-3g",
        redo_scf=True)

INFO:QMCTorch|
INFO:QMCTorch| SCF Calculation
INFO:QMCTorch|  Removing HLi_pyscf_sto-3g.hdf5 and redo SCF calculations
INFO:QMCTorch|  Running scf  calculation
converged SCF energy = -7.85928101642664
INFO:QMCTorch|  Molecule name       : HLi
INFO:QMCTorch|  Number of electrons : 4
INFO:QMCTorch|  SCF calculator      : pyscf
INFO:QMCTorch|  Basis set           : sto-3g
INFO:QMCTorch|  SCF                 : HF
INFO:QMCTorch|  Number of AOs       : 6
INFO:QMCTorch|  Number of MOs       : 6
INFO:QMCTorch|  SCF Energy          : -7.859 Hartree

We are going to use here to predefined Jastrow factors, one for electron-electron interactions and one for electon-nuclei interactions. Both use a fully connected neural network

# elec-elec jastrow factor
jastrow_ee = JastrowFactorElecElec(mol, FCEE)

# elec-nuclei jastrow factor
jastrow_en = JastrowFactorElecNuclei(mol, FCEN)

We can then pass a list containing these two jastrow factors to the wave function to combine them

wf = SlaterJastrow(mol, jastrow=[jastrow_ee, jastrow_en])

INFO:QMCTorch|
INFO:QMCTorch| Wave Function
INFO:QMCTorch|  Jastrow factor      : True
INFO:QMCTorch|  Jastrow kernel      : FullyConnectedJastrowKernel + FullyConnectedJastrowKernel
INFO:QMCTorch|  Highest MO included : 6
INFO:QMCTorch|  Configurations      : ground_state
INFO:QMCTorch|  Number of confs     : 1
INFO:QMCTorch|  Kinetic energy      : jacobi
INFO:QMCTorch|  Number var  param   : 367
INFO:QMCTorch|  Cuda support        : False

pos = torch.rand(10, wf.nelec*3)
print(wf(pos))

tensor([[-0.1107],
        [-0.1423],
        [-0.1496],
        [ 0.0150],
        [ 0.0299],
        [ 0.1291],
        [-0.1664],
        [ 0.1220],
        [ 0.2130],
        [-0.0605]], grad_fn=<MulBackward0>)