from qmctorch.ase import QMCTorch
from qmctorch.ase.optimizer import TorchOptimizer
from qmctorch.sampler.symmetry import Cinfv, Dinfh
from ase import Atoms
from ase.optimize import GoodOldQuasiNewton, FIRE
from ase.io import write
import torch
import numpy as np
from qmctorch.utils.plot_data import plot_walkers_traj, plot_correlation_coefficient

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

torch.random.manual_seed(0)
np.random.seed(0)

Define the ASE molecule and calculator

We first define the ASE molecule as ususal for ASE

lih = Atoms('LiH', positions=[(0, 0, 0), (0, 0, 3.14)])

We then attach a QMCTorch calculator to the atom object

lih.calc = QMCTorch()

Most of the options can directly be configured through the ASE calculator, via namespaces. The SCF options:

# SCF options
lih.calc.scf_options.calculator = 'adf'
lih.calc.scf_options.basis = 'dzp'

The options of the wave function

# WF options
lih.calc.wf_options.configs = 'single_double(2,4)'
lih.calc.wf_options.mix_mo = False
lih.calc.wf_options.orthogonalize_mo = False
lih.calc.wf_options.jastrow.kernel_kwargs = {'w':1.0}

The sampler options

# sampler options
lih.calc.sampler_options.nwalkers = 1000
lih.calc.sampler_options.nstep  = 500
lih.calc.sampler_options.step_size = 0.5
lih.calc.sampler_options.ntherm = -1
lih.calc.sampler_options.ndecor = 10
lih.calc.sampler_options.symmetry = None

The solver options:

# solver options
lih.calc.solver_options.freeze = []
lih.calc.solver_options.niter = 5
lih.calc.solver_options.tqdm = True
lih.calc.solver_options.grad = 'manual'

The resampling options

# options for the resampling
lih.calc.solver_options.resampling.mode = 'update'
lih.calc.solver_options.resampling.resample_every = 1
lih.calc.solver_options.resampling.ntherm_update = 100

We then need to initialize the calculator with all these options

# Optimize the wave function
lih.calc.initialize()

INFO:QMCTorch|
INFO:QMCTorch| SCF Calculation
INFO:QMCTorch|  Running scf  calculation
[15.04|15:03:07] PLAMS working folder: /home/nico/QMCTorch/docs/notebooks/plams_workdir
INFO:QMCTorch|  Molecule name       : HLi
INFO:QMCTorch|  Number of electrons : 4
INFO:QMCTorch|  SCF calculator      : adf
INFO:QMCTorch|  Basis set           : dzp
INFO:QMCTorch|  SCF                 : HF
INFO:QMCTorch|  Number of AOs       : 18
INFO:QMCTorch|  Number of MOs       : 17
INFO:QMCTorch|  SCF Energy          : -7.907 Hartree
INFO:QMCTorch|
INFO:QMCTorch| Wave Function
INFO:QMCTorch|  Backflow            : False
INFO:QMCTorch|  Jastrow factor      : True
INFO:QMCTorch|  Jastrow kernel      : ee -> PadeJastrowKernel
INFO:QMCTorch|  Highest MO included : 17
INFO:QMCTorch|  Configurations      : single_double(2,4)
INFO:QMCTorch|  Number of confs     : 16
INFO:QMCTorch|  Kinetic energy      : jacobi
INFO:QMCTorch|  Number var  param   : 347
INFO:QMCTorch|  Cuda support        : False
INFO:QMCTorch|
INFO:QMCTorch| Monte-Carlo Sampler
INFO:QMCTorch|  Number of walkers   : 1000
INFO:QMCTorch|  Number of steps     : 500
INFO:QMCTorch|  Step size           : 0.5
INFO:QMCTorch|  Thermalization steps: -1
INFO:QMCTorch|  Decorelation steps  : 10
INFO:QMCTorch|  Walkers init pos    : atomic
INFO:QMCTorch|  Move type           : all-elec
INFO:QMCTorch|  Move proba          : normal
INFO:QMCTorch|
INFO:QMCTorch| QMC Solver
INFO:QMCTorch|  WaveFunction        : SlaterJastrow
INFO:QMCTorch|  Sampler             : Metropolis
INFO:QMCTorch|  Optimizer           : Adam

As for any ASE calculator the energy of the system can be calculated as following

lih.get_potential_energy()

INFO:QMCTorch|  Initial Sampling    :

INFO:QMCTorch|  Sampling: 100%|██████████| 500/500 [00:41<00:00, 12.19it/s]

INFO:QMCTorch|   Acceptance rate     : 18.60 %
INFO:QMCTorch|   Timing statistics   : 12.19 steps/sec.
INFO:QMCTorch|   Total Time          : 41.03 sec.
INFO:QMCTorch|  done in 42.20 sec.
INFO:QMCTorch|
INFO:QMCTorch|  Optimization
INFO:QMCTorch|  Task                :
INFO:QMCTorch|  Number Parameters   : 341
INFO:QMCTorch|  Number of epoch     : 5
INFO:QMCTorch|  Batch size          : 1000
INFO:QMCTorch|  Loss function       : energy
INFO:QMCTorch|  Clip Loss           : False
INFO:QMCTorch|  Gradients           : manual
INFO:QMCTorch|  Resampling mode     : update
INFO:QMCTorch|  Resampling every    : 1
INFO:QMCTorch|  Resampling steps    : 50
INFO:QMCTorch|  Output file         : HLi_adf_dzp_QMCTorch.hdf5
INFO:QMCTorch|  Checkpoint every    : None
INFO:QMCTorch|
INFO:QMCTorch|
INFO:QMCTorch|  epoch 0 | 1000 sampling points
INFO:QMCTorch|  energy   : -7.710885 +/- 0.027782
INFO:QMCTorch|  variance : 0.878548
INFO:QMCTorch|  epoch done in 3.14 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 1 | 1000 sampling points
INFO:QMCTorch|  energy   : -7.731887 +/- 0.026397
INFO:QMCTorch|  variance : 0.834741
INFO:QMCTorch|  epoch done in 3.13 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 2 | 1000 sampling points
INFO:QMCTorch|  energy   : -7.727705 +/- 0.027536
INFO:QMCTorch|  variance : 0.870760
INFO:QMCTorch|  epoch done in 2.55 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 3 | 1000 sampling points
INFO:QMCTorch|  energy   : -7.739452 +/- 0.027085
INFO:QMCTorch|  variance : 0.856514
INFO:QMCTorch|  epoch done in 2.62 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 4 | 1000 sampling points
INFO:QMCTorch|  energy   : -7.715939 +/- 0.035746
INFO:QMCTorch|  variance : 1.130375
INFO:QMCTorch|  epoch done in 1.94 sec.
INFO:QMCTorch|
INFO:QMCTorch|  Single Point Calculation : 1000 walkers | 500 steps

INFO:QMCTorch|  Sampling: 100%|██████████| 500/500 [00:25<00:00, 19.62it/s]

INFO:QMCTorch|   Acceptance rate     : 18.02 %
INFO:QMCTorch|   Timing statistics   : 19.62 steps/sec.
INFO:QMCTorch|   Total Time          : 25.49 sec.
INFO:QMCTorch|  Energy   : -7.905802 +/- 0.019444
INFO:QMCTorch|  Variance : 0.378056
INFO:QMCTorch|  Size     : 1000

tensor(-7.9058)