{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Combining several Jastrow Factors\n",
    "\n",
    "It is often useful to use mutliple jastrow factors to go beyond the simple electron-electron Jastrow. We show here how to do that easily through our "
   ]
  },
  {
   "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 torch\n",
    "from qmctorch.scf import Molecule\n",
    "from qmctorch.wavefunction import SlaterJastrow\n",
    "\n",
    "from qmctorch.wavefunction.jastrows.elec_elec import (\n",
    "    JastrowFactor as JastrowFactorElecElec,\n",
    "    FullyConnectedJastrowKernel as FCEE,\n",
    ")\n",
    "from qmctorch.wavefunction.jastrows.elec_nuclei import (\n",
    "    JastrowFactor as JastrowFactorElecNuclei,\n",
    "    FullyConnectedJastrowKernel as FCEN,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use this base class to create a new Jastrow Factor. This is done in the same way one would create\n",
    "a new neural network layer in pytorch."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first need a molecule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| SCF Calculation\n",
      "INFO:QMCTorch|  Running scf  calculation\n",
      "converged SCF energy = -7.85928101642664\n",
      "INFO:QMCTorch|  Molecule name       : HLi\n",
      "INFO:QMCTorch|  Number of electrons : 4\n",
      "INFO:QMCTorch|  SCF calculator      : pyscf\n",
      "INFO:QMCTorch|  Basis set           : sto-3g\n",
      "INFO:QMCTorch|  SCF                 : HF\n",
      "INFO:QMCTorch|  Number of AOs       : 6\n",
      "INFO:QMCTorch|  Number of MOs       : 6\n",
      "INFO:QMCTorch|  SCF Energy          : -7.859 Hartree\n"
     ]
    }
   ],
   "source": [
    "mol = Molecule(\n",
    "        atom=\"Li 0 0 0; H 0 0 3.14\", \n",
    "        unit='bohr', \n",
    "        calculator=\"pyscf\",\n",
    "        basis=\"sto-3g\",\n",
    "        redo_scf=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to use here to predefined Jastrow factors, one for electron-electron interactions and one for electon-nuclei interactions. Both use a fully connected neural network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# elec-elec jastrow factor\n",
    "jastrow_ee = JastrowFactorElecElec(mol, FCEE)\n",
    "\n",
    "# elec-nuclei jastrow factor\n",
    "jastrow_en = JastrowFactorElecNuclei(mol, FCEN)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then pass a list containing these two jastrow factors to the wave function to combine them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| Wave Function\n",
      "INFO:QMCTorch|  Backflow            : False\n",
      "INFO:QMCTorch|  Jastrow factor      : True\n",
      "INFO:QMCTorch|  Jastrow kernel      : FullyConnectedJastrowKernel + FullyConnectedJastrowKernel\n",
      "INFO:QMCTorch|  Highest MO included : 6\n",
      "INFO:QMCTorch|  Configurations      : ground_state\n",
      "INFO:QMCTorch|  Number of confs     : 1\n",
      "INFO:QMCTorch|  Kinetic energy      : jacobi\n",
      "INFO:QMCTorch|  Number var  param   : 367\n",
      "INFO:QMCTorch|  Cuda support        : False\n"
     ]
    }
   ],
   "source": [
    "wf = SlaterJastrow(mol, jastrow=[jastrow_ee, jastrow_en])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0935],\n",
      "        [-0.3280],\n",
      "        [-1.4404],\n",
      "        [ 0.1365],\n",
      "        [ 2.3277],\n",
      "        [ 0.0043],\n",
      "        [ 0.0639],\n",
      "        [-0.0180],\n",
      "        [ 0.5653],\n",
      "        [-0.0892]], grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "pos = torch.rand(10, wf.nelec*3)\n",
    "print(wf(pos))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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
}