{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|  ____    __  ______________             _\n",
      "INFO:QMCTorch| / __ \\  /  |/  / ___/_  __/__  ________/ /  \n",
      "INFO:QMCTorch|/ /_/ / / /|_/ / /__  / / / _ \\/ __/ __/ _ \\ \n",
      "INFO:QMCTorch|\\___\\_\\/_/  /_/\\___/ /_/  \\___/_/  \\__/_//_/ \n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|0.4.0\n"
     ]
    }
   ],
   "source": [
    "from qmctorch.ase import QMCTorch     \n",
    "from qmctorch.ase.optimizer import TorchOptimizer\n",
    "from qmctorch.sampler.symmetry import Cinfv, Dinfh\n",
    "from ase import Atoms \n",
    "from ase.optimize import GoodOldQuasiNewton, FIRE\n",
    "from ase.io import write\n",
    "import torch\n",
    "import numpy as np\n",
    "from qmctorch.utils.plot_data import plot_walkers_traj, plot_correlation_coefficient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.random.manual_seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the ASE molecule and calculator\n",
    "\n",
    "We first define the ASE molecule as ususal for ASE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "lih = Atoms('LiH', positions=[(0, 0, 0), (0, 0, 3.14)])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then attach a QMCTorch calculator to the atom object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "lih.calc = QMCTorch()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most of the options can directly be configured through the ASE calculator, via namespaces. The SCF options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SCF options\n",
    "lih.calc.scf_options.calculator = 'adf'\n",
    "lih.calc.scf_options.basis = 'dzp'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The options of the wave function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WF options\n",
    "lih.calc.wf_options.configs = 'single_double(2,4)'\n",
    "lih.calc.wf_options.mix_mo = False\n",
    "lih.calc.wf_options.orthogonalize_mo = False\n",
    "lih.calc.wf_options.jastrow.kernel_kwargs = {'w':1.0}\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The sampler options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sampler options\n",
    "lih.calc.sampler_options.nwalkers = 1000\n",
    "lih.calc.sampler_options.nstep  = 500\n",
    "lih.calc.sampler_options.step_size = 0.5\n",
    "lih.calc.sampler_options.ntherm = -1\n",
    "lih.calc.sampler_options.ndecor = 10\n",
    "lih.calc.sampler_options.symmetry = None"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The solver options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solver options\n",
    "lih.calc.solver_options.freeze = []\n",
    "lih.calc.solver_options.niter = 5\n",
    "lih.calc.solver_options.tqdm = True\n",
    "lih.calc.solver_options.grad = 'manual'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resampling options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# options for the resampling\n",
    "lih.calc.solver_options.resampling.mode = 'update'\n",
    "lih.calc.solver_options.resampling.resample_every = 1\n",
    "lih.calc.solver_options.resampling.ntherm_update = 100"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then need to initialize the calculator with all these options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| SCF Calculation\n",
      "INFO:QMCTorch|  Running scf  calculation\n",
      "[15.04|15:03:07] PLAMS working folder: /home/nico/QMCTorch/docs/notebooks/plams_workdir\n",
      "INFO:QMCTorch|  Molecule name       : HLi\n",
      "INFO:QMCTorch|  Number of electrons : 4\n",
      "INFO:QMCTorch|  SCF calculator      : adf\n",
      "INFO:QMCTorch|  Basis set           : dzp\n",
      "INFO:QMCTorch|  SCF                 : HF\n",
      "INFO:QMCTorch|  Number of AOs       : 18\n",
      "INFO:QMCTorch|  Number of MOs       : 17\n",
      "INFO:QMCTorch|  SCF Energy          : -7.907 Hartree\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| Wave Function\n",
      "INFO:QMCTorch|  Backflow            : False\n",
      "INFO:QMCTorch|  Jastrow factor      : True\n",
      "INFO:QMCTorch|  Jastrow kernel      : ee -> PadeJastrowKernel\n",
      "INFO:QMCTorch|  Highest MO included : 17\n",
      "INFO:QMCTorch|  Configurations      : single_double(2,4)\n",
      "INFO:QMCTorch|  Number of confs     : 16\n",
      "INFO:QMCTorch|  Kinetic energy      : jacobi\n",
      "INFO:QMCTorch|  Number var  param   : 347\n",
      "INFO:QMCTorch|  Cuda support        : False\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| Monte-Carlo Sampler\n",
      "INFO:QMCTorch|  Number of walkers   : 1000\n",
      "INFO:QMCTorch|  Number of steps     : 500\n",
      "INFO:QMCTorch|  Step size           : 0.5\n",
      "INFO:QMCTorch|  Thermalization steps: -1\n",
      "INFO:QMCTorch|  Decorelation steps  : 10\n",
      "INFO:QMCTorch|  Walkers init pos    : atomic\n",
      "INFO:QMCTorch|  Move type           : all-elec\n",
      "INFO:QMCTorch|  Move proba          : normal\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| QMC Solver \n",
      "INFO:QMCTorch|  WaveFunction        : SlaterJastrow\n",
      "INFO:QMCTorch|  Sampler             : Metropolis\n",
      "INFO:QMCTorch|  Optimizer           : Adam\n"
     ]
    }
   ],
   "source": [
    "# Optimize the wave function\n",
    "lih.calc.initialize()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As for any ASE calculator the energy of the system can be calculated as following"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|  Initial Sampling    :\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|  Sampling: 100%|██████████| 500/500 [00:41<00:00, 12.19it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|   Acceptance rate     : 18.60 %\n",
      "INFO:QMCTorch|   Timing statistics   : 12.19 steps/sec.\n",
      "INFO:QMCTorch|   Total Time          : 41.03 sec.\n",
      "INFO:QMCTorch|  done in 42.20 sec.\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|  Optimization\n",
      "INFO:QMCTorch|  Task                :\n",
      "INFO:QMCTorch|  Number Parameters   : 341\n",
      "INFO:QMCTorch|  Number of epoch     : 5\n",
      "INFO:QMCTorch|  Batch size          : 1000\n",
      "INFO:QMCTorch|  Loss function       : energy\n",
      "INFO:QMCTorch|  Clip Loss           : False\n",
      "INFO:QMCTorch|  Gradients           : manual\n",
      "INFO:QMCTorch|  Resampling mode     : update\n",
      "INFO:QMCTorch|  Resampling every    : 1\n",
      "INFO:QMCTorch|  Resampling steps    : 50\n",
      "INFO:QMCTorch|  Output file         : HLi_adf_dzp_QMCTorch.hdf5\n",
      "INFO:QMCTorch|  Checkpoint every    : None\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|  epoch 0 | 1000 sampling points\n",
      "INFO:QMCTorch|  energy   : -7.710885 +/- 0.027782\n",
      "INFO:QMCTorch|  variance : 0.878548\n",
      "INFO:QMCTorch|  epoch done in 3.14 sec.\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|  epoch 1 | 1000 sampling points\n",
      "INFO:QMCTorch|  energy   : -7.731887 +/- 0.026397\n",
      "INFO:QMCTorch|  variance : 0.834741\n",
      "INFO:QMCTorch|  epoch done in 3.13 sec.\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|  epoch 2 | 1000 sampling points\n",
      "INFO:QMCTorch|  energy   : -7.727705 +/- 0.027536\n",
      "INFO:QMCTorch|  variance : 0.870760\n",
      "INFO:QMCTorch|  epoch done in 2.55 sec.\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|  epoch 3 | 1000 sampling points\n",
      "INFO:QMCTorch|  energy   : -7.739452 +/- 0.027085\n",
      "INFO:QMCTorch|  variance : 0.856514\n",
      "INFO:QMCTorch|  epoch done in 2.62 sec.\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|  epoch 4 | 1000 sampling points\n",
      "INFO:QMCTorch|  energy   : -7.715939 +/- 0.035746\n",
      "INFO:QMCTorch|  variance : 1.130375\n",
      "INFO:QMCTorch|  epoch done in 1.94 sec.\n",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch|  Single Point Calculation : 1000 walkers | 500 steps\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|  Sampling: 100%|██████████| 500/500 [00:25<00:00, 19.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|   Acceptance rate     : 18.02 %\n",
      "INFO:QMCTorch|   Timing statistics   : 19.62 steps/sec.\n",
      "INFO:QMCTorch|   Total Time          : 25.49 sec.\n",
      "INFO:QMCTorch|  Energy   : -7.905802 +/- 0.019444\n",
      "INFO:QMCTorch|  Variance : 0.378056\n",
      "INFO:QMCTorch|  Size     : 1000\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor(-7.9058)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lih.get_potential_energy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qmctorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}