{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "# Creating your own Jastrow Factor\n", "\n", "We present here how to create your own electron-electron Jastrow factor and use it in QMCTorch.\n", "During the import you must import the base class of the electron-electron Jastrow. 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" ] } ], "source": [ "import torch\n", "from qmctorch.scf import Molecule\n", "from qmctorch.wavefunction import SlaterJastrow\n", "from qmctorch.wavefunction.jastrows.elec_elec.jastrow_factor_electron_electron import JastrowFactorElectronElectron\n", "from qmctorch.wavefunction.jastrows.elec_elec.kernels import JastrowKernelElectronElectronBase" ] }, { "cell_type": "code", "execution_count": 2, "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": [ "mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', calculator='pyscf', unit='bohr', redo_scf=True)" ] }, { "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." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "from torch import nn \n", "class MyJastrowKernel(JastrowKernelElectronElectronBase):\n", " def __init__(self, nup, ndown, cuda, size=16):\n", " super().__init__(nup, ndown, cuda)\n", " self.fc1 = nn.Linear(1, size, bias=False)\n", " self.fc2 = nn.Linear(size, 1, bias=False)\n", " def forward(self, x):\n", " nbatch, npair = x.shape\n", " x = x.reshape(-1,1)\n", " x = self.fc2(self.fc1(x))\n", " return x.reshape(nbatch, npair)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "As seen above the prototype of the class constructor must be: \n", "\n", "`def __init__(self, nup, ndown, cuda, **kwargs)`\n", "\n", "The list of keyword argument can contain any pairs such as ``size=16``.\n", "\n", "\n", "This Jastrow use two fully connected layers. The size of the hidden layer is here controlled by a keyword argument ``size`` whose defauilt value is 16\n", "It is important to note that the calculation of the first and second derivative of the jastrow kernel wrt the electronic positions are then done via automatic differentiation as implemented in the `JastrowKernelElectronElectronBase` class. Hence there is no need to derive and implement these derivatives. However it is necessary that the ``forward`` function, which takes as input a ``torch.tensor`` of\n", "dimension ``[Nbatch, Npair]`` first reshape this tensor to ``[Nbatch*Npair,1]``, then applies the transformation on this tensor and finally reshape\n", "the output tensor to ``[Nbatch, Npair]``.\n", "\n", "To use this new Jastrow kernel in the `SlaterJastrow` wave function ansatz we first need to instantiate a Jastrow factor that uses the kernel. " ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "jastrow = JastrowFactorElectronElectron(mol, MyJastrowKernel, kernel_kwargs={'size': 64})" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "This jastrow factor can then be passed as an argument of the `SlaterJastrow` wavefunction." ] }, { "cell_type": "code", "execution_count": 5, "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 : ee -> MyJastrowKernel\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 : 249\n", "INFO:QMCTorch| Cuda support : False\n" ] } ], "source": [ "wf = SlaterJastrow(mol, jastrow=jastrow)" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "tensor([[0.0767],\n", " [0.1802],\n", " [0.1636],\n", " [0.3930],\n", " [0.2627],\n", " [0.1758],\n", " [0.1102],\n", " [0.1478],\n", " [0.1833],\n", " [0.1580]], grad_fn=)\n" ] } ], "source": [ "pos = torch.rand(10, wf.nelec*3)\n", "print(wf(pos))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "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 }