diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 9402c20f..767cd31b 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -10,7 +10,7 @@ jobs: strategy: fail-fast: false matrix: - version: [3.8] + version: [3.8,3.9] steps: - name: Cancel Previous Runs diff --git a/docs/example/single_point/h2.py b/docs/example/single_point/h2.py index 27017364..3a1c603c 100644 --- a/docs/example/single_point/h2.py +++ b/docs/example/single_point/h2.py @@ -15,7 +15,7 @@ # define the wave function wf = SlaterJastrow(mol, kinetic='jacobi', - configs='ground_state', jastrow=jastrow).gto2sto() + configs='ground_state', jastrow=jastrow) #.gto2sto() # sampler sampler = Metropolis(nwalkers=1000, nstep=1000, step_size=0.25, 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 diff --git a/docs/source/modules.rst b/docs/source/modules.rst index 762c26c8..31572603 100644 --- a/docs/source/modules.rst +++ b/docs/source/modules.rst @@ -1,9 +1,10 @@ -Subpackages ------------ +Python Interface +------------------- .. toctree:: - :maxdepth: 4 + :maxdepth: 1 + :hidden: qmctorch.sampler qmctorch.scf @@ -11,7 +12,7 @@ Subpackages qmctorch.utils qmctorch.wavefunction -Module contents +Reference --------------- .. automodule:: qmctorch diff --git a/qmctorch/wavefunction/jastrows/elec_elec_nuclei/kernels/boys_handy_jastrow_kernel.py b/qmctorch/wavefunction/jastrows/elec_elec_nuclei/kernels/boys_handy_jastrow_kernel.py index eae30937..0ec75e3f 100644 --- a/qmctorch/wavefunction/jastrows/elec_elec_nuclei/kernels/boys_handy_jastrow_kernel.py +++ b/qmctorch/wavefunction/jastrows/elec_elec_nuclei/kernels/boys_handy_jastrow_kernel.py @@ -9,7 +9,7 @@ class BoysHandyJastrowKernel(JastrowKernelElectronElectronNucleiBase): def __init__( self, nup, ndown, atomic_pos, cuda, nterm=5 ): # pylint: disable=too-many-arguments - """Defines a Boys Handy jastrow factors. + r"""Defines a Boys Handy jastrow factors. J.W. Moskowitz et. al Correlated Monte Carlo Wave Functions for Some Cations and Anions of the First Row Atoms diff --git a/qmctorch/wavefunction/orbitals/norm_orbital.py b/qmctorch/wavefunction/orbitals/norm_orbital.py index 8b7d5546..fd73b567 100644 --- a/qmctorch/wavefunction/orbitals/norm_orbital.py +++ b/qmctorch/wavefunction/orbitals/norm_orbital.py @@ -1,6 +1,6 @@ import torch import numpy as np - +from scipy.special import factorial2 def atomic_orbital_norm(basis): """Computes the norm of the atomic orbitals @@ -75,7 +75,6 @@ def norm_gaussian_spherical(bas_n, bas_exp): torch.tensor: normalization factor """ - from scipy.special import factorial2 as f2 bas_n = torch.tensor(bas_n) bas_n = bas_n + 1.0 @@ -104,7 +103,6 @@ def norm_slater_cartesian(a, b, c, n, exp): Returns: torch.tensor: normalization factor """ - from scipy.special import factorial2 as f2 lvals = a + b + c + n + 1.0 @@ -141,16 +139,24 @@ def norm_gaussian_cartesian(a, b, c, exp): torch.tensor: normalization factor """ - from scipy.special import factorial2 as f2 - pref = torch.as_tensor((2 * exp / np.pi) ** (0.75)) am1 = (2 * a - 1).astype("int") x = (4 * exp) ** (a / 2) / torch.sqrt(torch.as_tensor(f2(am1))) bm1 = (2 * b - 1).astype("int") y = (4 * exp) ** (b / 2) / torch.sqrt(torch.as_tensor(f2(bm1))) - + cm1 = (2 * c - 1).astype("int") - z = (4 * exp) ** (c / 2) / torch.sqrt(torch.as_tensor(f2(cm1))) + z = (4 * exp) ** (c / 2) / torch.sqrt(torch.as_tensor(f2(cm1))) return (pref * x * y * z).type(torch.get_default_dtype()) + +def f2(x): + """Returns the f2 of x with f2(x<1) = 1 as implemented in scipy 1.10. + """ + # compute the x!! + out = factorial2(x) + + # set all the elements lower than 1 to 1 + out[out<1] = 1 + return out \ No newline at end of file diff --git a/qmctorch/wavefunction/slater_jastrow.py b/qmctorch/wavefunction/slater_jastrow.py index 94e6a30a..6f919e2f 100644 --- a/qmctorch/wavefunction/slater_jastrow.py +++ b/qmctorch/wavefunction/slater_jastrow.py @@ -341,7 +341,7 @@ def gradients_jacobi(self, x, sum_grad=False, pdf=False): are computed following .. math:: - \\nabla \\Psi(R) = \\left( \\nabla J(R) \\right) \\Sigma + J(R) \\left(\\nabla \Sigma \\right) + \\nabla \\Psi(R) = \\left( \\nabla J(R) \\right) \\Sigma + J(R) \\left(\\nabla \\Sigma \\right) with @@ -453,11 +453,11 @@ def get_kinetic_operator(self, x, ao, dao, d2ao, mo): return -0.5 * bkin def kinetic_energy_jacobi_backflow(self, x, **kwargs): - r"""Compute the value of the kinetic enery using the Jacobi Formula. + """Compute the value of the kinetic enery using the Jacobi Formula. .. math:: - \\frac{\Delta (J(R) \Psi(R))}{ J(R) \Psi(R)} = \\frac{\\Delta J(R)}{J(R} + \\frac{\\Delta (J(R) \\Psi(R))}{ J(R) \\Psi(R)} = \\frac{\\Delta J(R)}{J(R} + 2 \\frac{\\nabla J(R)}{J(R)} \\frac{\\nabla \\Psi(R)}{\\Psi(R)} + \\frac{\\Delta \\Psi(R)}{\\Psi(R)}