Geometry Optimization

We present here a complete example on how to use QMCTorch to optimize the geometry of a H2 molecule. As previously the first task is to import all the modules needed

[1]:
from torch import optim
from torch.optim import Adam
from qmctorch.wavefunction import SlaterJastrow
from qmctorch.solver import Solver
from qmctorch.sampler import Metropolis
from qmctorch.scf import Molecule
from qmctorch.utils.plot_data import plot_energy
from qmctorch.utils import set_torch_double_precision
set_torch_double_precision()
INFO:QMCTorch|  ____    __  ______________             _
INFO:QMCTorch| / __ \  /  |/  / ___/_  __/__  ________/ /
INFO:QMCTorch|/ /_/ / / /|_/ / /__  / / / _ \/ __/ __/ _ \
INFO:QMCTorch|\___\_\/_/  /_/\___/ /_/  \___/_/  \__/_//_/

We then define the molecule. We create a H2 molecule with a short atomic distance that we will then try to relax

[2]:
mol = Molecule(atom = 'H 0. 0. -0.5; H 0. 0. 0.5', unit='bohr',
               calculator='pyscf', basis='sto-3g', 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

As for the wave function optimization example we need to define a wave function, a sampler an optimizer and a solver.

[3]:
# wave function with only the ground state determinant
wf = SlaterJastrow(mol, configs='single_double(2,2)').gto2sto()

# sampler
sampler = Metropolis(nwalkers=1000, nstep=200, step_size=0.5,
                     nelec=wf.nelec, ndim=wf.ndim,
                     ntherm=-1, ndecor=100,
                     init=mol.domain('normal'),
                     move={'type': 'all-elec', 'proba': 'normal'})
# optimizer
opt = Adam(wf.parameters(), lr=0.005)

# scheduler
scheduler = optim.lr_scheduler.StepLR(opt, step_size=20, gamma=0.75)
INFO:QMCTorch|
INFO:QMCTorch| Wave Function
INFO:QMCTorch|  Jastrow factor      : True
INFO:QMCTorch|  Jastrow kernel      : ee -> PadeJastrowKernel
INFO:QMCTorch|  Highest MO included : 2
INFO:QMCTorch|  Configurations      : single_double(2,2)
INFO:QMCTorch|  Number of confs     : 4
INFO:QMCTorch|  Kinetic energy      : jacobi
INFO:QMCTorch|  Number var  param   : 21
INFO:QMCTorch|  Cuda support        : False
INFO:QMCTorch|  Fit GTOs to STOs  :
INFO:QMCTorch|
INFO:QMCTorch| Wave Function
INFO:QMCTorch|  Jastrow factor      : True
INFO:QMCTorch|  Jastrow kernel      : ee -> PadeJastrowKernel
INFO:QMCTorch|  Highest MO included : 2
INFO:QMCTorch|  Configurations      : single_double(2,2)
INFO:QMCTorch|  Number of confs     : 4
INFO:QMCTorch|  Kinetic energy      : jacobi
INFO:QMCTorch|  Number var  param   : 17
INFO:QMCTorch|  Cuda support        : False
INFO:QMCTorch|
INFO:QMCTorch| Monte-Carlo Sampler
INFO:QMCTorch|  Number of walkers   : 1000
INFO:QMCTorch|  Number of steps     : 200
INFO:QMCTorch|  Step size           : 0.5
INFO:QMCTorch|  Thermalization steps: -1
INFO:QMCTorch|  Decorelation steps  : 100
INFO:QMCTorch|  Walkers init pos    : normal
INFO:QMCTorch|  Move type           : all-elec
INFO:QMCTorch|  Move proba          : normal
[4]:
solver = Solver(wf=wf,
                       sampler=sampler,
                       optimizer=opt,
                       scheduler=None)
solver.configure(loss='energy', grad='auto', track=['local_energy','geometry'], freeze = ['ao'])
INFO:QMCTorch|
INFO:QMCTorch| QMC Solver
INFO:QMCTorch|  WaveFunction        : SlaterJastrow
INFO:QMCTorch|  Sampler             : Metropolis
INFO:QMCTorch|  Optimizer           : Adam

Note that we must set grad='auto' to perform a geometry optimization. To optimize the geometry of the molecule we must use the set_params_requires_grad method of the solver. We here set all the wave function parameters and all the atomic positions as variational parameters, leading to a simultaneous optimization of the wave function and of the atomic coordinates. We can then run the optimization here using 50 epochs

[5]:
solver.set_params_requires_grad(wf_params=False, geo_params=True)
obs = solver.run(50)

INFO:QMCTorch|
INFO:QMCTorch|  Optimization
INFO:QMCTorch|  Task                :
INFO:QMCTorch|  Number Parameters   : 6
INFO:QMCTorch|  Number of epoch     : 50
INFO:QMCTorch|  Batch size          : 1000
INFO:QMCTorch|  Loss function       : energy
INFO:QMCTorch|  Clip Loss           : False
INFO:QMCTorch|  Gradients           : auto
INFO:QMCTorch|  Resampling mode     : update
INFO:QMCTorch|  Resampling every    : 1
INFO:QMCTorch|  Resampling steps    : 25
INFO:QMCTorch|  Output file         : H2_pyscf_sto-3g_QMCTorch.hdf5
INFO:QMCTorch|  Checkpoint every    : None
INFO:QMCTorch|
INFO:QMCTorch|
INFO:QMCTorch|  epoch 0 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.051299 +/- 0.011930
INFO:QMCTorch|  variance : 0.377259
INFO:QMCTorch|  epoch done in 0.08 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 1 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.039677 +/- 0.011319
INFO:QMCTorch|  variance : 0.357927
INFO:QMCTorch|  epoch done in 0.11 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 2 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.083498 +/- 0.011489
INFO:QMCTorch|  variance : 0.363315
INFO:QMCTorch|  epoch done in 0.13 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 3 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.096306 +/- 0.012919
INFO:QMCTorch|  variance : 0.408541
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 4 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.073275 +/- 0.011114
INFO:QMCTorch|  variance : 0.351444
INFO:QMCTorch|  epoch done in 0.14 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 5 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.091809 +/- 0.011652
INFO:QMCTorch|  variance : 0.368468
INFO:QMCTorch|  epoch done in 0.11 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 6 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.086389 +/- 0.011696
INFO:QMCTorch|  variance : 0.369854
INFO:QMCTorch|  epoch done in 0.51 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 7 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.106363 +/- 0.015236
INFO:QMCTorch|  variance : 0.481804
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 8 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.103029 +/- 0.011551
INFO:QMCTorch|  variance : 0.365266
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 9 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.104694 +/- 0.010543
INFO:QMCTorch|  variance : 0.333408
INFO:QMCTorch|  epoch done in 0.13 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 10 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.101184 +/- 0.010772
INFO:QMCTorch|  variance : 0.340643
INFO:QMCTorch|  epoch done in 0.83 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 11 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.118775 +/- 0.010548
INFO:QMCTorch|  variance : 0.333553
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 12 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.111163 +/- 0.010536
INFO:QMCTorch|  variance : 0.333193
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 13 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.093243 +/- 0.010991
INFO:QMCTorch|  variance : 0.347571
INFO:QMCTorch|  epoch done in 0.15 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 14 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.098148 +/- 0.009786
INFO:QMCTorch|  variance : 0.309451
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 15 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.110943 +/- 0.009883
INFO:QMCTorch|  variance : 0.312525
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 16 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.130522 +/- 0.010770
INFO:QMCTorch|  variance : 0.340589
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 17 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.119832 +/- 0.010397
INFO:QMCTorch|  variance : 0.328774
INFO:QMCTorch|  epoch done in 0.84 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 18 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.114476 +/- 0.010044
INFO:QMCTorch|  variance : 0.317618
INFO:QMCTorch|  epoch done in 0.11 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 19 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.127139 +/- 0.009782
INFO:QMCTorch|  variance : 0.309344
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 20 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.141108 +/- 0.010158
INFO:QMCTorch|  variance : 0.321225
INFO:QMCTorch|  epoch done in 0.16 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 21 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.120979 +/- 0.009740
INFO:QMCTorch|  variance : 0.307994
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 22 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.130151 +/- 0.010171
INFO:QMCTorch|  variance : 0.321641
INFO:QMCTorch|  epoch done in 0.16 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 23 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.126904 +/- 0.009565
INFO:QMCTorch|  variance : 0.302478
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 24 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.140780 +/- 0.009760
INFO:QMCTorch|  variance : 0.308645
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 25 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.122666 +/- 0.009864
INFO:QMCTorch|  variance : 0.311916
INFO:QMCTorch|  epoch done in 0.14 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 26 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.124710 +/- 0.009316
INFO:QMCTorch|  variance : 0.294589
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 27 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.130382 +/- 0.009266
INFO:QMCTorch|  variance : 0.293020
INFO:QMCTorch|  epoch done in 0.18 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 28 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.124134 +/- 0.009203
INFO:QMCTorch|  variance : 0.291026
INFO:QMCTorch|  epoch done in 0.11 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 29 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.148873 +/- 0.008958
INFO:QMCTorch|  variance : 0.283281
INFO:QMCTorch|  epoch done in 0.11 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 30 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.151410 +/- 0.009124
INFO:QMCTorch|  variance : 0.288521
INFO:QMCTorch|  epoch done in 0.11 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 31 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.132741 +/- 0.008910
INFO:QMCTorch|  variance : 0.281774
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 32 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.157418 +/- 0.009251
INFO:QMCTorch|  variance : 0.292552
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 33 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.139773 +/- 0.009447
INFO:QMCTorch|  variance : 0.298734
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 34 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.141193 +/- 0.009174
INFO:QMCTorch|  variance : 0.290105
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 35 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.142530 +/- 0.011028
INFO:QMCTorch|  variance : 0.348741
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 36 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.135991 +/- 0.009584
INFO:QMCTorch|  variance : 0.303085
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 37 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.149069 +/- 0.009036
INFO:QMCTorch|  variance : 0.285733
INFO:QMCTorch|  epoch done in 0.11 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 38 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.146994 +/- 0.009392
INFO:QMCTorch|  variance : 0.297009
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 39 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.141153 +/- 0.008568
INFO:QMCTorch|  variance : 0.270943
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 40 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.146822 +/- 0.008714
INFO:QMCTorch|  variance : 0.275570
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 41 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.137690 +/- 0.008836
INFO:QMCTorch|  variance : 0.279429
INFO:QMCTorch|  epoch done in 0.10 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 42 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.140852 +/- 0.008685
INFO:QMCTorch|  variance : 0.274655
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 43 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.141535 +/- 0.008415
INFO:QMCTorch|  variance : 0.266091
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 44 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.156989 +/- 0.008510
INFO:QMCTorch|  variance : 0.269122
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 45 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.136263 +/- 0.008505
INFO:QMCTorch|  variance : 0.268963
INFO:QMCTorch|  epoch done in 0.09 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 46 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.135993 +/- 0.009332
INFO:QMCTorch|  variance : 0.295104
INFO:QMCTorch|  epoch done in 0.15 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 47 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.151198 +/- 0.008536
INFO:QMCTorch|  variance : 0.269921
INFO:QMCTorch|  epoch done in 0.12 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 48 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.142816 +/- 0.008155
INFO:QMCTorch|  variance : 0.257869
INFO:QMCTorch|  epoch done in 0.13 sec.
INFO:QMCTorch|
INFO:QMCTorch|  epoch 49 | 1000 sampling points
INFO:QMCTorch|  energy   : -1.140517 +/- 0.008837
INFO:QMCTorch|  variance : 0.279447
INFO:QMCTorch|  epoch done in 0.12 sec.
[6]:
solver.save_traj('h2_traj.xyz', obs)
[7]:
plot_energy(obs.local_energy)
../_images/notebooks_geoopt_10_0.png