Add shls_slice kwarg for pyscf.df.incore (fix pyscf#1729)

sunqm · Aug 16, 2024 · 48cc4f2 · 48cc4f2
1 parent df09359
commit 48cc4f2
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/pyscf/df/incore.py b/pyscf/df/incore.py
@@ -38,22 +38,34 @@
 
 
 def aux_e2(mol, auxmol_or_auxbasis, intor='int3c2e', aosym='s1', comp=None, out=None,
-           cintopt=None):
+           cintopt=None, shls_slice=None):
     '''3-center AO integrals (ij|L), where L is the auxiliary basis.
 
     Kwargs:
-        cintopt : Libcint-3.14 and newer version support to compute int3c2e
-            without the opt for the 3rd index.  It can be precomputed to
-            reduce the overhead of cintopt initialization repeatedly.
+        cintopt :
+            Precomputing certain pair-shell data. It can be created by
 
             cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, 'int3c2e')
+
+        shls_slice : 6-element tuple
+            Label the start-stop shells for each index in the integral tensor.
+            For the (ij|aux) = intor('int3c2e'), the tuple should be given as
+            (ish_start, ish_end, jsh_start, jsh_end, aux_start, aux_end)
     '''
     if isinstance(auxmol_or_auxbasis, gto.MoleBase):
         auxmol = auxmol_or_auxbasis
     else:
         auxbasis = auxmol_or_auxbasis
         auxmol = addons.make_auxmol(mol, auxbasis)
-    shls_slice = (0, mol.nbas, 0, mol.nbas, mol.nbas, mol.nbas+auxmol.nbas)
+    if shls_slice is None:
+        shls_slice = (0, mol.nbas, 0, mol.nbas,
+                      mol.nbas, mol.nbas+auxmol.nbas)
+    else:
+        assert len(shls_slice) == 6
+        assert shls_slice[5] < auxmol.nbas
+        shls_slice = list(shls_slice)
+        shls_slice[4] += mol.nbas
+        shls_slice[5] += mol.nbas
 
     # Extract the call of the two lines below
     #  pmol = gto.mole.conc_mol(mol, auxmol)

diff --git a/pyscf/gto/mole.py b/pyscf/gto/mole.py
@@ -3441,6 +3441,11 @@ def intor(self, intor, comp=None, hermi=0, aosym='s1', out=None,
                 | 1 : hermitian
                 | 2 : anti-hermitian
 
+            shls_slice : 4-element, 6-element or 8-element tuple
+                Label the start-stop shells for each index in the integral.
+                For example, the 8-element tuple for the 2-electron integral
+                tensor (ij|kl) = intor('int2e') are specified as
+                (ish_start, ish_end, jsh_start, jsh_end, ksh_start, ksh_end, lsh_start, lsh_end)
             grids : ndarray
                 Coordinates of grids for the int1e_grids integrals