From cc397de88cbc40f5f8d86501d435ab72379551e1 Mon Sep 17 00:00:00 2001 From: Nico Date: Tue, 5 Dec 2023 14:12:41 +0100 Subject: [PATCH] updated rst in the doc --- docs/rst/qmctorch.rst | 269 ++++++++++++++++++++++++------------------ 1 file changed, 155 insertions(+), 114 deletions(-) diff --git a/docs/rst/qmctorch.rst b/docs/rst/qmctorch.rst index 1f81fa67..80c7d8b3 100644 --- a/docs/rst/qmctorch.rst +++ b/docs/rst/qmctorch.rst @@ -1,4 +1,4 @@ -Wave Function ansatz in QMCTorch +Wave Function Ansatz in QMCTorch =========================================== `QMCTorch` allows to epxress the wave function ususally used by QMC practitioner as neural network. The most generic architecture of the @@ -15,22 +15,28 @@ These atomic orbital values are then transformed to molecular orbital values thr Then a Slater determinant layer extract the different determinants contained in the wave function. Users can there as well specify wich determinants they require. The weighted sum of the determinants is then computed and finally muliplied with the value of the Jastrow factor. -Different wave function forms have been implemented to easily create and use wave function ansatz. These different functional forms differ mainly by the Jastrow factor they use and the presence of backflow transformation or not. +The main wave function in QMCTorch is implemented in the ``SlaterJastrow`` class. The definition of the class is as follows : -Two-body Jastrow factors -^^^^^^^^^^^^^^^^^^^^^^^^^^ -In its simplest form the Jastrow factor only depends on the electron-electron distances. This means that the Jastrow layer only has a single kernel function :math:`K_{ee}`. -This Jastrow factor can be applied globally, or different Jastrow factors can be applied to individual orbitals. In addition a Backflow transformation can be added or not to the definition -of the wave function. We therefore have the following wave function available: +.. code-block:: python + + class SlaterJastrow(WaveFunction): + def __init__( + self, + mol, + jastrow='default', + backflow=None, + configs="ground_state", + kinetic="jacobi", + cuda=False, + include_all_mo=True, + ): -* ``SlaterJastrow``: A simple wave function containing an electron-electron Jastrow factor and a sum of Slater determinants -* ``SlaterOrbitalDependentJastrow``: A ``SlaterJastrow`` for but each molecular orbitals has its own Jastrow factor -* ``SlaterJastrowBackflow``: A ``SlaterJastrow`` wave function with backflow transformation for the electrons +Different functional form can be created from this class depending on the need of the user. We review here a few of these forms. -Slater Jastrow Wave Function ----------------------------------------- +Simple Slater Jastrow Wave Function +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ The simplest wave function implemented in `QMCTorch` is a Slater Jastrow form. Through a series of transformations the Slater Jastrow function computes: @@ -62,7 +68,15 @@ The determinantal parts in the expression of :math:`\Psi` are given by the spin- A ``SlaterJastrow`` wave function can instantiated following : ->>> wf = SlaterJastrow(mol, configs='single_double(2,2)', jastrow_kernel=PadeJastrowKernel) +.. code-block:: python + + from qmctorch.scf import Molecule + from qmctorch.wavefunction.slater_jastrow import SlaterJastrow + from qmctorch.wavefunction.jastrows.elec_elec.jastrow_factor_electron_electron import JastrowFactorElectronElectron + from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel + mol = Molecule('H 0 0 0; H 0 0 1') + jastrow = JastrowFactorElectronElectron(mol, PadeJastrowKernel) + wf = SlaterJastrow(mol, configs='single_double(2,2)', jastrow=jastrow) The ``SlaterJastrow`` takes as first mandiatory argument a ``Molecule`` instance. The Slater determinants required in the calculation are specified with the ``configs`` arguments which can take the following values : @@ -72,135 +86,162 @@ are specified with the ``configs`` arguments which can take the following values * ``configs='single(n,m)'`` : only single excitation using n electron and m orbitals * ``configs='single_double(n,m)'`` : only single/double excitation using n electron and m orbitals -Finally the kernel function of the Jastrow factor can be specifed using the ``jastrow_kernel`` -The ``SlaterJastrow`` class accepts other initialisation arguments to fine tune some advanced settings. The default values -of these arguments are adequeate for most cases. - -Orbital dependent Slater Jastrow Wave Function ---------------------------------------------------- - -A slight modification of the the Slater Jastrow is obtained by making the the Jastrow factor can be made orbital dependent. -This is implemented in the ``SlaterOrbitalDependentJastrow`` that can be instantiated as: - ->>> from qmctorch.wavefunction import SlaterOrbitalDependentJastrow ->>> from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel ->>> wf = SlaterOrbitalDependentJastrow(mol, configs='single_double(2,4)' ->>> jastrow_kernel=PadeJastrowKernel) - -Slater Jastrow Backflow Wave Function ----------------------------------------- - -The Slater Jastrow Backflow wave function builds on the the Slater Jastrow wavefunction but adds a backflow transformation to -the electronic positions. Following this transformation, each electron becomes a quasi-particle whose position depends on all -electronic positions. The backflow transformation is given by : - -.. math:: - - q(x_i) = x_i + \sum_{j\neq i} \text{Kernel}(r_{ij}) (x_i-x_j) - -The kernel of the transformation can be any function that depends on the distance between two electrons. A popular kernel -is simply the inverse function : - -.. math:: - \text{Kernel}(r_{ij}) = \frac{\omega}{r_{ij}} - -and is the default value in QMCTorch. However any other kernel function can be implemented and used in the code. +Finally the Jastrow factor can be specifed using the ``jastrow``. We used here a Pade-Jastrow kernel that is already implemented in QMCTorch -The wave function is then constructed as : - -.. math:: +Custom Jastrow factor +^^^^^^^^^^^^^^^^^^^^^^^^^^ - \Psi(R) = J(R) \sum_n c_n D_n^{\uparrow}(Q) D_n^{\downarrow}(Q) +It is possible to define custom Jastrow factor and use these forms in the definition of the wave function. -The Jastrow factor is still computed using the original positions of the electrons while the determinant part uses the -backflow transformed positions. One can define such wave function with: +.. code-block:: python ->>> from qmctorch.wavefunction import SlaterJastrowBackFlow ->>> from qmctorch.wavefunction.orbitals.backflow.kernels import BackFlowKernelInverse ->>> from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel ->>> ->>> wf = SlaterJastrowBackFlow(mol, ->>> configs='single_double(2,2)', ->>> jastrow_kernel=PadeJastrowKernel, ->>> backflow_kernel=BackFlowKernelInverse) + from torch import nn + from qmctorch.wavefunction import SlaterJastrow + from qmctorch.wavefunction.jastrows.elec_elec.jastrow_factor_electron_electron import JastrowFactorElectronElectron + from qmctorch.wavefunction.jastrows.elec_elec.kernels import JastrowKernelElectronElectronBase -Compared to the ``SlaterJastrow`` wave function, the kernel of the backflow transformation must be specified. By default the inverse kernel will be used. + class MyJastrowKernel(JastrowKernelElectronElectronBase): + def __init__(self, nup, ndown, cuda, size=16): + super().__init__(nup, ndown, cuda) + self.fc1 = nn.Linear(1, size, bias=False) + self.fc2 = nn.Linear(size, 1, bias=False) + def forward(self, x): + nbatch, npair = x.shape + x = x.reshape(-1,1) + x = self.fc2(self.fc1(x)) + return x.reshape(nbatch, npair) -Orbital Dependent Backflow Transformation -****************************************** + mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', calculator='pyscf', unit='bohr', redo_scf=True) -The backflow transformation can be different for each atomic orbitals. + jastrow = JastrowFactorElectronElectron(mol, MyJastrowKernel, kernel_kwargs={'size': 64}) -.. math:: + wf = SlaterJastrow(mol, jastrow=jastrow) - q^\alpha(x_i) = x_i + \sum_{j\neq i} \text{Kernel}^\alpha(r_{ij}) (x_i-x_j) -where each orbital has its dedicated backflow kernel. This provides much more flexibility when optimizing the wave function. +Combining Several Jastrow Factors +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -This wave function can be used with +As shown on the figure above it is possible to combine several Jastrow factors to account for not only the electron-electron correlations but also electron-nuclei and three body terms. +This can easily be done by passing a list of Jastrow factors to the `SlaterJastrow` wave function. +For example if we want to combine a fully connected electron-electron neural Jastrow factor with a fully connected electron-nuclei neural Jastrow, we can simply use: ->>> from qmctorch.wavefunction import SlaterJastrowBackFlow ->>> from qmctorch.wavefunction.orbitals.backflow.kernels import BackFlowKernelInverse ->>> from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel ->>> ->>> wf = SlaterJastrowBackFlow(mol, ->>> configs='single_double(2,2)', ->>> jastrow_kernel=PadeJastrowKernel, ->>> orbital_dependent_backflow=True, ->>> backflow_kernel=BackFlowKernelInverse) +.. code-block:: python + import torch + from qmctorch.scf import Molecule + from qmctorch.wavefunction import SlaterJastrow -Many-Body Jastrow factors -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + from qmctorch.wavefunction.jastrows.elec_elec import ( + JastrowFactor as JastrowFactorElecElec, + FullyConnectedJastrowKernel as FCEE, + ) + from qmctorch.wavefunction.jastrows.elec_nuclei import ( + JastrowFactor as JastrowFactorElecNuclei, + FullyConnectedJastrowKernel as FCEN, + ) -Jastrow factors can also depends on the electron-nuclei distances and the many body terms involving two electrons and one nuclei. -In that case the Jastrow factor depends on all the kernel function represented in the figure above. A backflow transformation can also be added to the definition of the wave function. -As a result we have the following wave function forms available. + mol = Molecule( + atom="Li 0 0 0; H 0 0 3.14", + unit='bohr', + calculator="pyscf", + basis="sto-3g", + redo_scf=True) -* ``SlaterManyBodyJastrow``: A wave function that contains a many body Jastrow factor and a sum of Slater determinants with backflow transformation for the electrons -* ``SlaterManyBodyJastrowBackflow``: A ``SlaterManyBodyJastrow`` wave function with a backflow transformation + jastrow_ee = JastrowFactorElecElec(mol, FCEE) + jastrow_en = JastrowFactorElecNuclei(mol, FCEN) + wf = SlaterJastrow(mol, jastrow=[jastrow_ee, jastrow_en]) -Many-Body Jastrow Wave Function ----------------------------------------- +Wave Functions with Backflow Transformations +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -The Jastrow factor combines here multiple terms that represent electron-electron, electron-nuclei and electron-electron-nuclei terms. +As seen on the figure above, a backflow transformation of the electronic positions can be added to the definition of the wave function. +Following this transformation, each electron becomes a quasi-particle whose position depends on all +electronic positions. The backflow transformation is given by : .. math:: - J(R_{at},r) = \exp\left( \sum_{i>> from qmctorch.wavefunction import SlaterManyBodyJastrow ->>> from qmctorch.wavefunction.jastrows.elec_elec.kernels import PadeJastrowKernel as PadeJastrowElecElec ->>> from qmctorch.wavefunction.jastrows.elec_nuclei.kernels import PadeJastrowKernel as PadeJastrowKernelElecNuc ->>> from qmctorch.wavefunction.jastrows.elec_elec_nuclei.kernels import BoysHandyJastrowKernel ->>> ->>> wf = SlaterManyBodyJastrow(mol, ->>> configs='single_double(2,2)', ->>> jastrow_kernel={ ->>> 'ee': PadeJastrowKernelElecElec, ->>> 'en': PadeJastrowKernelElecNuc, ->>> 'een': BoysHandyJastrowKernel}) +.. math:: + \text{Kernel}(r_{ij}) = \frac{\omega}{r_{ij}} +and is the default value in QMCTorch. However any other kernel function can be implemented and used in the code. +The wave function is then constructed as : -Many-Body Jastrow Wave Function with backflow transformation ------------------------------------------------------------------- +.. math:: -A backflow transformation can be used together with the many body Jastrow + \Psi(R) = J(R) \sum_n c_n D_n^{\uparrow}(Q) D_n^{\downarrow}(Q) +The Jastrow factor is still computed using the original positions of the electrons while the determinant part uses the +backflow transformed positions. One can define such wave function with: ->>> from qmctorch.wavefunction import SlaterManyBodyJastrowBackflow ->>> from qmctorch.wavefunction.jastrows.elec_elec.kernels.pade_jastrow_kernel import PadeJastrowKernel as PadeJastrowElecElec ->>> from qmctorch.wavefunction.jastrows.elec_nuclei.kernels.pade_jastrow_kernel import PadeJastrowKernel as PadeJastrowKernelElecNuc ->>> from qmctorch.wavefunction.jastrows.elec_elec_nuclei.kernels.boys_handy_jastrow_kernel import BoysHandyJastrowKernel ->>> ->>> wf = SlaterManyBodyJastrowBackflow(mol, ->>> configs='single_double(2,2)', ->>> jastrow_kernel={ ->>> 'ee': PadeJastrowKernelElecElec, ->>> 'en': PadeJastrowKernelElecNuc, ->>> 'een': BoysHandyJastrowKernel}, ->>> backflow_kernel=BackFlowKernelInverse) +.. code-block:: python + + from qmctorch.scf import Molecule + from qmctorch.wavefunction.slater_jastrow import SlaterJastrow + + from qmctorch.wavefunction.jastrows.elec_elec import JastrowFactor, PadeJastrowKernel + + from qmctorch.wavefunction.orbitals.backflow import ( + BackFlowTransformation, + BackFlowKernelInverse, + ) + + # molecule + mol = Molecule( + atom="Li 0 0 0; H 0 0 3.015", + unit="bohr", + calculator="pyscf", + basis="sto-3g", + redo_scf=True, + ) + + # define jastrow factor + jastrow = JastrowFactor(mol, PadeJastrowKernel) + + # define backflow trans + backflow = BackFlowTransformation(mol, BackFlowKernelInverse) + + # define the wave function + wf = SlaterJastrow( + mol, + configs="single_double(2,2)", + jastrow=jastrow, + backflow=backflow, + ) + +Custom Backflow Transformation +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +As for the Jastrow factor, it is possible to create custom backlfow transformations and use them in the definition of the wave function. +For example to define a fully connected backflow kernel and use it we can use: + +.. code-block:: python + + import torch + from torch import nn + from qmctorch.scf import Molecule + from qmctorch.wavefunction import SlaterJastrow + from qmctorch.wavefunction.orbitals.backflow.kernels import BackFlowKernelBase + from qmctorch.wavefunction.orbitals.backflow import BackFlowTransformation + + class MyBackflowKernel(BackFlowKernelBase): + def __init__(self, mol, cuda, size=16): + super().__init__(mol, cuda) + self.fc1 = nn.Linear(1, size, bias=False) + self.fc2 = nn.Linear(size, 1, bias=False) + def forward(self, x): + original_shape = x.shape + x = x.reshape(-1,1) + x = self.fc2(self.fc1(x)) + return x.reshape(*original_shape) + + mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', unit='bohr', redo_scf=True) + backflow = BackFlowTransformation(mol, MyBackflowKernel, backflow_kernel_kwargs={'size': 8}) + wf = SlaterJastrow(mol, backflow=backflow) \ No newline at end of file