{
 "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": [
    "import qmctorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Jastrow Factor\n",
    "\n",
    "The wave function of the molecular system is written as :\n",
    "\n",
    "$$\n",
    "\\Psi(R) = J(R) \\sum_n c_n \\det(A_\\uparrow(r_\\uparrow)) \\det(A_\\downarrow(r_\\downarrow))\n",
    "$$\n",
    "\n",
    "where $J(R)$ is the so called Jastrow factor, and $A_\\uparrow$($A_\\downarrow$) is the matrix of the molecular orbitals values for the spin up(down) electron\n",
    "\n",
    "The Jastrow factor is written as the exponential of a kernel function :\n",
    "\n",
    "$$\n",
    "J(R) = \\exp\\left( \\sum_{i<j}  f(r_{ij}) \\right) \n",
    "$$\n",
    "\n",
    "where $r_{ij}$ denotes the distance between electrons $i$ and $j$. The kernel function $f(r_{ij})$ can take differet forms. Traditionally it is written as a Pade-Jastrow function\n",
    "\n",
    "$$\n",
    "f(r_{ij}) = \\frac{a r_{ij}}{1+ \\omega r_{ij}}\n",
    "$$\n",
    "\n",
    "where $a$ is a fixed weight and $\\omega$ a variational parameter."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The electron-electron Jastrow factor class (`JastrowFactorElectronElectron`) orchestrates the calculations of the Jastrow factors (and its 1st and 2nd derivatives). It can take different kernel functions as for example the `PadeJastrowKernel`. The example below shows how to use this class.\n",
    "\n",
    "In the simple Pade-Jastrow function the value of the $a$ parameter is given by the so-called electron-electron cusp conditions. These conditions are here to insured that the value of the total energy remains finite at the coalescence point, i.e. when two electrons have the same positions. \n",
    "\n",
    "It can be shown that this is respected when\n",
    "\n",
    "$$\n",
    "\\frac{1}{J(R)}{\\frac{\\partial J(R)}{\\partial r_{ij}} |_{r_{ij}=0}} = \\frac{1}{4} \n",
    "$$\n",
    "\n",
    "for same-spin electrons and \n",
    "\n",
    "$$\n",
    "\\frac{1}{J(R)}{\\frac{\\partial J(R)}{\\partial r_{ij}} |_{r_{ij}=0}} = \\frac{1}{2} \n",
    "$$\n",
    "\n",
    "for opposite spin electrons.\n",
    "\n",
    "For the Pade-Jastrow function this translates into $a=1/4$ ($a=1/2$) for same(opposite) spin electrons. However realizing the same condition for the fully connected Jastrow factor is still to be clarified. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| SCF Calculation\n",
      "INFO:QMCTorch|  Removing H2_pyscf_dzp.hdf5 and redo SCF calculations\n",
      "INFO:QMCTorch|  Running scf  calculation\n",
      "converged SCF energy = -1.07280585930373\n",
      "INFO:QMCTorch|  Molecule name       : H2\n",
      "INFO:QMCTorch|  Number of electrons : 2\n",
      "INFO:QMCTorch|  SCF calculator      : pyscf\n",
      "INFO:QMCTorch|  Basis set           : dzp\n",
      "INFO:QMCTorch|  SCF                 : HF\n",
      "INFO:QMCTorch|  Number of AOs       : 10\n",
      "INFO:QMCTorch|  Number of MOs       : 10\n",
      "INFO:QMCTorch|  SCF Energy          : -1.073 Hartree\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from qmctorch.scf import Molecule\n",
    "from qmctorch.wavefunction.jastrows.elec_elec.jastrow_factor_electron_electron import JastrowFactorElectronElectron\n",
    "from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel\n",
    "\n",
    "# create the molecule\n",
    "mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', calculator='pyscf', unit='bohr', redo_scf=True)\n",
    "\n",
    "# define the jastrow factor\n",
    "jastrow = JastrowFactorElectronElectron(\n",
    "            mol,\n",
    "            PadeJastrowKernel,\n",
    "            kernel_kwargs={'w': 0.1})\n",
    "\n",
    "# define random electronic positions\n",
    "nbatch = 10\n",
    "pos = torch.rand(nbatch, mol.nelec * 3)\n",
    "\n",
    "# compute the jastrow\n",
    "jval = jastrow(pos)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Jastrows\n",
    "\n",
    "The functional form $f(r_{ij}) = ar_{ij}(1+\\omega r_{ij})^{-1}$ only offers one single variational parameter and does not offer a lot of flexibility. It is however possible to replace that form by a simple fully connected neural network\n",
    "\n",
    "This network takes a single input value ($r_{ij}$) and outputs a single value, i.e. the value of the kernel. A simple 2 layer fully connected neurakl network  Jastrow Kernel has been implemented in the `FullyConnectedJastrowKernel` class that can be used in the `JastrowFactorElectronElectron` as follow:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| SCF Calculation\n",
      "INFO:QMCTorch|  Removing H2_pyscf_dzp.hdf5 and redo SCF calculations\n",
      "INFO:QMCTorch|  Running scf  calculation\n",
      "converged SCF energy = -1.07280585930373\n",
      "INFO:QMCTorch|  Molecule name       : H2\n",
      "INFO:QMCTorch|  Number of electrons : 2\n",
      "INFO:QMCTorch|  SCF calculator      : pyscf\n",
      "INFO:QMCTorch|  Basis set           : dzp\n",
      "INFO:QMCTorch|  SCF                 : HF\n",
      "INFO:QMCTorch|  Number of AOs       : 10\n",
      "INFO:QMCTorch|  Number of MOs       : 10\n",
      "INFO:QMCTorch|  SCF Energy          : -1.073 Hartree\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from qmctorch.scf import Molecule\n",
    "from qmctorch.wavefunction.jastrows.elec_elec.jastrow_factor_electron_electron import JastrowFactorElectronElectron\n",
    "from qmctorch.wavefunction.jastrows.elec_elec.kernels import FullyConnectedJastrowKernel\n",
    "\n",
    "# number of spin up/down electrons\n",
    "mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', calculator='pyscf', unit='bohr', redo_scf=True)\n",
    "\n",
    "# define the jastrow factor\n",
    "jastrow = JastrowFactorElectronElectron(\n",
    "            mol,\n",
    "            FullyConnectedJastrowKernel,\n",
    "            kernel_kwargs={'size1': 32, 'size2': 64})\n",
    "\n",
    "# define random electronic positions\n",
    "nbatch = 10\n",
    "pos = torch.rand(nbatch, mol.nelec * 3)\n",
    "\n",
    "# compute the jastrow\n",
    "jval = jastrow(pos)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SlaterJastrow Wave function\n",
    "\n",
    "Both Jastrow factors can be used to define a `SlaterJastrow` wavefunction. The example below shows how to do that for a `LiH` molecule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| SCF Calculation\n",
      "INFO:QMCTorch|  Reusing scf results from LiH_pyscf_sto-3g.hdf5\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 -> FullyConnectedJastrowKernel\n",
      "INFO:QMCTorch|  Highest MO included : 6\n",
      "INFO:QMCTorch|  Configurations      : single_double(2,2)\n",
      "INFO:QMCTorch|  Number of confs     : 4\n",
      "INFO:QMCTorch|  Kinetic energy      : jacobi\n",
      "INFO:QMCTorch|  Number var  param   : 2210\n",
      "INFO:QMCTorch|  Cuda support        : False\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from qmctorch.scf import Molecule\n",
    "from qmctorch.wavefunction import SlaterJastrow\n",
    "from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel, FullyConnectedJastrowKernel\n",
    "\n",
    "# define the molecule\n",
    "mol = Molecule(\n",
    "            atom='Li 0 0 0; H 0 0 3.14',\n",
    "            unit='bohr',\n",
    "            calculator='pyscf',\n",
    "            basis='sto-3g')\n",
    "\n",
    "\n",
    "# define the jastrow factor\n",
    "jastrow = JastrowFactorElectronElectron(\n",
    "            mol,\n",
    "            FullyConnectedJastrowKernel,\n",
    "            kernel_kwargs={'size1': 32, 'size2': 64})\n",
    "\n",
    "# define the Slater Jastrow wavefunction\n",
    "wf = SlaterJastrow(mol,\n",
    "                   jastrow=jastrow,\n",
    "                   configs='single_double(2,2)')\n",
    "\n",
    "# define random electronic positions\n",
    "nbatch = 10\n",
    "pos = torch.rand(nbatch, mol.nelec * 3)\n",
    "\n",
    "# compute the value of the wave function\n",
    "wfval = wf(pos)"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}