{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating your own Backflow transformation\n",
    "\n",
    "We present here how to create your own backflow transformation and use it in QMCTorch.\n",
    "During the import you must import the base class of the backflow kernel. We aso create a H2 molecule"
   ]
  },
  {
   "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",
      "INFO:QMCTorch|\n",
      "INFO:QMCTorch| SCF Calculation\n",
      "INFO:QMCTorch|  Removing H2_adf_dzp.hdf5 and redo SCF calculations\n",
      "INFO:QMCTorch|  Running scf  calculation\n",
      "[15.04|14:38:27] PLAMS working folder: /home/nico/QMCTorch/docs/notebooks/plams_workdir\n",
      "INFO:QMCTorch|  Molecule name       : H2\n",
      "INFO:QMCTorch|  Number of electrons : 2\n",
      "INFO:QMCTorch|  SCF calculator      : adf\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.082 Hartree\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from qmctorch.scf import Molecule\n",
    "from qmctorch.wavefunction import SlaterJastrow\n",
    "from qmctorch.wavefunction.orbitals.backflow.kernels import BackFlowKernelBase\n",
    "from qmctorch.wavefunction.orbitals.backflow import BackFlowTransformation\n",
    "\n",
    "mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', unit='bohr', redo_scf=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use this base class to create a new backflow transformation kernel.\n",
    "This is done in the same way one would create a new neural network layer in pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn \n",
    "class MyBackflowKernel(BackFlowKernelBase):\n",
    "    def __init__(self, mol, cuda, size=16):\n",
    "        super().__init__(mol, cuda)\n",
    "        self.fc1 = nn.Linear(1, size, bias=False)\n",
    "        self.fc2 = nn.Linear(size, 1, bias=False)\n",
    "    def _backflow_kernel(self, x):\n",
    "        original_shape = x.shape\n",
    "        x = x.reshape(-1,1)\n",
    "        x = self.fc2(self.fc1(x))\n",
    "        return x.reshape(*original_shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This backflow transformation consists of two fully connected layers. The calculation of the first and second derivative are then done via automatic differentiation as implemented in the `BackFlowKernelBase` class. To use this new kernel in the `SlaterJastrow` wave function ansatz we first need to instantiate a backflow layer using this kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "backflow = BackFlowTransformation(mol, MyBackflowKernel, backflow_kernel_kwargs={'size': 8})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use this backflow transformation in the call of the wave function:"
   ]
  },
  {
   "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            : True\n",
      "INFO:QMCTorch|  Backflow kernel     : MyBackflowKernel\n",
      "INFO:QMCTorch|  Jastrow factor      : True\n",
      "INFO:QMCTorch|  Jastrow kernel      : ee -> PadeJastrowKernel\n",
      "INFO:QMCTorch|  Highest MO included : 10\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   : 134\n",
      "INFO:QMCTorch|  Cuda support        : False\n"
     ]
    }
   ],
   "source": [
    "wf = SlaterJastrow(mol, backflow=backflow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2102],\n",
      "        [0.0593],\n",
      "        [0.0715],\n",
      "        [0.0618],\n",
      "        [0.0581],\n",
      "        [0.1396],\n",
      "        [0.1200],\n",
      "        [0.0711],\n",
      "        [0.1123],\n",
      "        [0.1754]], 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
}