Creating your own Backflow transformationΒΆ

We present here how to create your own backflow transformation and use it in QMCTorch. During the import you must import the base class of the backflow kernel. We aso create a H2 molecule

[4]:
import torch
from qmctorch.scf import Molecule
from qmctorch.wavefunction import SlaterJastrowBackFlow
from qmctorch.wavefunction.orbitals.backflow.kernels import BackFlowKernelBase
mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', unit='bohr', redo_scf=True)
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

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

[5]:
from torch import nn
class MyBackflow(BackFlowKernelBase):
    def __init__(self, mol, cuda, size=16):
        super().__init__(mol, cuda)
        self.fc1 = nn.Linear(1, size, bias=False)
        self.fc2 = nn.Linear(size, 1, bias=False)
    def forward(self, x):
        original_shape = x.shape
        x = x.reshape(-1,1)
        x = self.fc2(self.fc1(x))
        return x.reshape(*original_shape)

This backflow transformation consists of two fully connected layers. The calculation of the first and second derivative are then done via automatic differentiation as implemented in the BackFlowKernelBase class. To use this new kernel in the SlaterJastrowBackFlow wave function ansatz we simply pass the class name as argument of the backflow_kernel keyword argument :

[6]:
wf = SlaterJastrowBackFlow(mol,
                   backflow_kernel=MyBackflow,
                   backflow_kernel_kwargs={'size' : 64})
INFO:QMCTorch|
INFO:QMCTorch| Wave Function
INFO:QMCTorch|  Jastrow factor      : True
INFO:QMCTorch|  Jastrow kernel      : PadeJastrowKernel
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   : 146
INFO:QMCTorch|  Cuda support        : False
[7]:
pos = torch.rand(10, wf.nelec*3)
print(wf(pos))
tensor([[0.1134],
        [0.1509],
        [0.1096],
        [0.1093],
        [0.2632],
        [0.1523],
        [0.1253],
        [0.1424],
        [0.1324],
        [0.0665]], grad_fn=<MulBackward0>)
[ ]: