Feat: Add on-the-fly density-fitted fragment ERI transform (#38)

* Feat: Add on-the-fly density-fitted fragment ERI transform

Co-authored-by: hongzhouye <hzyechem@gmail.com>

* Add test for on-the-fly DF ERI transformation

* Add note about HF-in-HF error for int direct DF

---------

Co-authored-by: hongzhouye <hzyechem@gmail.com>

molbe/be_var.py

@@ -3,3 +3,4 @@
 SOLVER_CALL = 0
 SCRATCH=''
 CREATE_SCRATCH_DIR=False
+INTEGRAL_TRANSFORM_MAX_MEMORY=50 # in GB

molbe/eri_onthefly.py

@@ -0,0 +1,92 @@
+# Author(s): Minsik Cho, Hong-Zhou Ye
+
+import numpy, os
+
+from pyscf import lib
+from pyscf.df.addons import make_auxmol
+from pyscf.gto.moleintor import getints3c, make_loc, make_cintopt
+from pyscf.gto import mole
+from pyscf.ao2mo.addons import restore
+from scipy.linalg import cholesky, solve_triangular
+from . import be_var
+
+def integral_direct_DF(mf, Fobjs, file_eri, auxbasis=None):
+    """ Calculate AO density-fitted 3-center integrals on-the-fly and transform to Schmidt space for given fragment objects
+
+    Parameters
+    ----------
+    mf : pyscf.scf.RHF
+        Mean-field object for the chemical system (typically BE.mf)
+    Fobjs : list of BE.Frags
+        List containing fragment objects (typically BE.Fobjs)
+        The MO coefficients are taken from Frags.TA and the transformed ERIs are stored in Frags.dname as h5py datasets.
+    file_eri : h5py.File
+        HDF5 file object to store the transformed fragment ERIs
+    auxbasis : string, optional
+        Auxiliary basis used for density fitting. If not provided, use pyscf's default choice for the basis set used to construct mf object; by default None
+    """
+
+    def calculate_pqL(aux_range):
+        """ Internal function to calculate the 3-center integrals for a given range of auxiliary indices
+
+        Parameters
+        ----------
+        aux_range : tuple of int
+            (start index, end index) of the auxiliary basis functions to calculate the 3-center integrals, i.e. (pq|L) with L ∈ [start, end) is returned
+        """
+        if be_var.PRINT_LEVEL > 10: print('Start calculating (μν|P) for range', aux_range, flush=True)
+        p0, p1 = aux_range
+        shls_slice = (0, mf.mol.nbas, 0, mf.mol.nbas, mf.mol.nbas + p0, mf.mol.nbas + p1)
+        ints = getints3c(mf.mol._add_suffix('int3c2e'), atm, bas, env, shls_slice, 1, 's1', ao_loc, cintopt, out=None)
+        if be_var.PRINT_LEVEL > 10: print('Finish calculating (μν|P) for range', aux_range, flush=True)
+        return ints
+
+    def block_step_size(nfrag, naux, nao):
+        """ Internal function to calculate the block step size for the 3-center integrals calculation
+
+        Parameters
+        ----------
+        nfrag : int
+            Number of fragments
+        naux : int
+            Number of auxiliary basis functions
+        nao : int
+            Number of atomic orbitals
+        """
+        return max(1, int(
+            be_var.INTEGRAL_TRANSFORM_MAX_MEMORY * 1e9 / 8 / nao / nao / naux / nfrag
+        )) #max(int(500*.24e6/8/nao),1)
+
+    if be_var.PRINT_LEVEL > 2:
+        print('Evaluating fragment ERIs on-the-fly using density fitting...', flush=True)
+        print('In this case, note that HF-in-HF error includes DF error on top of numerical error from embedding.', flush=True)
+    auxmol = make_auxmol(mf.mol, auxbasis = auxbasis)
+    j2c = auxmol.intor(mf.mol._add_suffix('int2c2e'), hermi=1) # (L|M)
+    low = cholesky(j2c, lower=True)
+    pqL_frag = [numpy.zeros((auxmol.nao, fragobj.nao, fragobj.nao)) for fragobj in Fobjs] # place to store fragment (pq|L)
+    end = 0
+    atm, bas, env = mole.conc_env(mf.mol._atm, mf.mol._bas, mf.mol._env, auxmol._atm, auxmol._bas, auxmol._env)
+    ao_loc = make_loc(bas, mf.mol._add_suffix('int3c2e')); cintopt = make_cintopt(atm, bas, env, mf.mol._add_suffix('int3c2e'))
+    blockranges = [(x, y) for x, y in lib.prange(0, auxmol.nbas, block_step_size(len(Fobjs), auxmol.nbas, mf.mol.nao))]
+    if be_var.PRINT_LEVEL > 4:
+        print('Aux Basis Block Info: ', blockranges, flush=True)
+
+    for idx, ints in enumerate(lib.map_with_prefetch(calculate_pqL, blockranges)):
+        if be_var.PRINT_LEVEL > 4: print('Calculating pq|L block #', idx, blockranges[idx], flush=True)
+        # Transform pq (AO) to fragment space (ij)
+        start = end; end += ints.shape[2]
+        for fragidx in range(len(Fobjs)):
+            if be_var.PRINT_LEVEL > 10: print('(μν|P) -> (ij|P) for frag #', fragidx, flush=True)
+            Lqp = numpy.transpose(ints, axes=(2,1,0))
+            Lqi = Lqp @ Fobjs[fragidx].TA
+            Liq = numpy.moveaxis(Lqi, 2, 1)
+            pqL_frag[fragidx][start:end, :, :] = Liq @ Fobjs[fragidx].TA
+    # Fit to get B_{ij}^{L}
+    for fragidx in range(len(Fobjs)):
+        if be_var.PRINT_LEVEL > 10: print('Fitting B_{ij}^{L} for frag #', fragidx, flush=True)
+        b = pqL_frag[fragidx].reshape(auxmol.nao, -1)
+        bb = solve_triangular(low, b, lower=True, overwrite_b=True, check_finite=False)
+        if be_var.PRINT_LEVEL > 10: print('Finished obtaining B_{ij}^{L} for frag #', fragidx, flush=True)
+        eri_nosym = bb.T @ bb
+        eri = restore('4', eri_nosym, Fobjs[fragidx].nao)
+        file_eri.create_dataset(Fobjs[fragidx].dname, data=eri)

molbe/mbe.py

@@ -55,7 +55,8 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5',
                  mo_energy = None, 
                  save_file='storebe.pk',hci_pt=False,
                  nproc=1, ompnum=4,
-                 hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None):
+                 hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None,
+                 integral_direct_DF=False, auxbasis = None):
         """
         Constructor for BE object.
 
@@ -85,6 +86,10 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5',
             Number of processors for parallel calculations, by default 1. If set to >1, threaded parallel computation is invoked.
         ompnum : int, optional
             Number of OpenMP threads, by default 4.
+        integral_direct_DF: bool, optional
+            If mf._eri is None (i.e. ERIs are not saved in memory using incore_anyway), this flag is used to determine if the ERIs are computed integral-directly using density fitting; by default False.
+        auxbasis : str, optional
+            Auxiliary basis for density fitting, by default None (uses default auxiliary basis defined in PySCF).
         """
 
         if restart:
@@ -111,6 +116,8 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5',
         self.unrestricted = False
         self.nproc = nproc
         self.ompnum = ompnum
+        self.integral_direct_DF = integral_direct_DF
+        self.auxbasis = auxbasis
 
         # Fragment information from fobj
         self.frag_type=fobj.frag_type
@@ -303,7 +310,8 @@ def initialize(self, eri_,compute_hf, restart=False):
             #   Yes -- ao2mo, incore version
             #   No  -- Do we have (ij|P) from density fitting?
             #            Yes -- ao2mo, outcore version, using saved (ij|P)
-            assert (not eri_ is None) or (hasattr(self.mf, 'with_df')), "Input mean-field object is missing ERI (mf._eri) or DF (mf.with_df) object. You may want to ensure that incore_anyway was set for non-DF calculations."
+            #            No  -- if integral_direct_DF is requested, invoke on-the-fly routine
+            assert (not eri_ is None) or (hasattr(self.mf, 'with_df')) or (self.integral_direct_DF), "Input mean-field object is missing ERI (mf._eri) or DF (mf.with_df) object AND integral direct DF routine was not requested. Please check your inputs."
             if not eri_ is None: # incore ao2mo using saved eri from mean-field calculation
                 for I in range(self.Nfrag):
                     eri = ao2mo.incore.full(eri_, self.Fobjs[I].TA, compact=True)
@@ -313,6 +321,15 @@ def initialize(self, eri_,compute_hf, restart=False):
                 for I in range(self.Nfrag):
                     eri = self.mf.with_df.ao2mo(self.Fobjs[I].TA, compact=True)
                     file_eri.create_dataset(self.Fobjs[I].dname, data=eri)
+            else:
+                # If ERIs are not saved on memory, compute fragment ERIs integral-direct
+                if self.integral_direct_DF: # Use density fitting to generate fragment ERIs on-the-fly
+                    from .eri_onthefly import integral_direct_DF
+                    integral_direct_DF(self.mf, self.Fobjs, file_eri, auxbasis=self.auxbasis)
+                else: # Calculate ERIs on-the-fly to generate fragment ERIs
+                    # TODO: Future feature to be implemented
+                    # NOTE: Ideally, we want AO shell pair screening for this.
+                    return NotImplementedError
         else:
             eri=None
 

tests/eri_onthefly_test.py

@@ -0,0 +1,28 @@
+"""
+This script tests the on-the-fly DF-based ERI transformation routine for molecular systems.
+
+Author(s): Minsik Cho
+"""
+
+import os, unittest
+from pyscf import gto, scf
+from molbe import fragpart, BE
+
+class TestDF_ontheflyERI(unittest.TestCase):
+    @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions")
+    def test_octane_BE2(self):
+        # Octane, cc-pvtz
+        mol = gto.M()
+        mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/octane.xyz')
+        mol.basis = 'cc-pvtz'
+        mol.charge = 0.; mol.spin = 0.
+        mol.build()
+        mf = scf.RHF(mol); mf.direct_scf = True; mf.kernel()
+        fobj = fragpart(frag_type='autogen', be_type='be2', mol = mol)
+        mybe = BE(mf, fobj, integral_direct_DF=True)
+        self.assertAlmostEqual(mybe.ebe_hf, mf.e_tot,
+                               msg = "HF-in-HF energy for Octane (BE2) does not match the HF energy!",
+                               delta = 1e-6)
+
+if __name__ == '__main__':
+    unittest.main()