Fully connected Network Jastrow

[NicoRenaud]

The fully connected Jastrow Factor can be used as :

from torch import optim
from torch import nn
from qmctorch.scf import Molecule
from qmctorch.wavefunction import SlaterJastrow
from qmctorch.solver import SolverSlaterJastrow
from qmctorch.sampler import Metropolis
from qmctorch.utils import set_torch_double_precision
from qmctorch.utils import (plot_energy, plot_data)
from jastrow import FullyConnectedJastrowKernel

set_torch_double_precision()

# define the molecule
mol = Molecule(atom='Li 0 0 0; H 0 0 3.015', 
               calculator='adf', basis='dzp', unit='bohr')

# define the wave function
wf = SlaterJastrow(mol, kinetic='jacobi',
                   configs='single_double(4,12)',
                   jastrow_kernel=FullyConnectedJastrowKernel,
                   cuda=True)

# sampler
sampler = Metropolis(nwalkers=10000,
                     nstep=2000, step_size=0.05,
                     ntherm=-1, ndecor=100,
                     nelec=wf.nelec, init=mol.domain('atomic'),
                     move={'type': 'all-elec', 'proba': 'normal'},
                    cuda=True)

# optimizer
lr_dict = [{'params': wf.jastrow.parameters(), 'lr': 1E-2},
           {'params': wf.ao.parameters(), 'lr': 1E-2},
           {'params': wf.mo.parameters(), 'lr': 1E-2},
           {'params': wf.fc.parameters(), 'lr': 1E-2}]
opt = optim.Adam(lr_dict, lr=1E-3)

# scheduler
scheduler = optim.lr_scheduler.StepLR(opt, step_size=500, gamma=0.90)

# QMC solver
solver = SolverSlaterJastrow(wf=wf, sampler=sampler, optimizer=opt, scheduler=scheduler)

# optimize the wave function
solver.track_observable(['local_energy','parameters'])

solver.configure_resampling(mode='update', resample_every=1, nstep_update=100)
solver.ortho_mo = False

#solver.configure(task='wf_opt',freeze=None)
obs = solver.run(500)

The architecture of the default FullyConnectedJastrowKernel is given by a simple three layer fc network with a Sigmoid activation function. (see qmctorch/wavefunction/jastrows/elec_elec/kernel/fully_connected_jastrow_kernel.py)

We should explore how different network architecture (number of layers, number of nodes in each layer, activation function, etc ...) affect the performance of the model