pyfock 0.2.0__tar.gz → 0.2.1__tar.gz

{pyfock-0.2.0/pyfock.egg-info → pyfock-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyfock
-Version: 0.2.0
+Version: 0.2.1
 Summary: A simplistic and efficient pure-python quantum chemistry library from Phys Whiz.
 Home-page: https://github.com/manassharma07/pyfock
 Download-URL: 

{pyfock-0.2.0 → pyfock-0.2.1}/pyfock/DFT.py

@@ -139,10 +139,12 @@ class DFT:
         Whether to use Rys quadrature for evaluating electron repulsion integrals.
 
     DF_algo : int
-        Algorithm selector for DF (reserved for developer use). 11 (default) is the shell-blocked
-        CPU/GPU algorithm that honours ``max_memory_ints3c2e`` and also provides RI-HF exchange
-        (``xc='HF'``, CPU); 12 adds multipole expansions for the far-field three-center integrals
-        (CPU, pure functionals; see ``multipole_options``); 10 is the previous default.
+        Algorithm selector for DF (reserved for developer use). 12 (default, CPU/GPU) is the
+        shell-blocked algorithm with multipole expansions for the far-field three-center integrals
+        (pure functionals; see ``multipole_options``); 11 is the same near-field algorithm without
+        the far field and is selected automatically for RI exact exchange (``xc='HF'`` and hybrid
+        functionals, CPU), which needs the complete three-center blocks; 10 is the former default.
+        Both 11 and 12 honour ``max_memory_ints3c2e``.
 
     multipole_options : dict
         Parameters of the far-field multipole expansions of DF_algo=12 (``precision``, ``lmax``,
@@ -150,8 +152,10 @@ class DFT:
         :data:`pyfock.Integrals.df_algo12_helpers.DEFAULT_OPTIONS`.
 
     max_memory_ints3c2e : float or None
-        Memory budget (GB) for the stored three-center integrals with DF_algo=11
-        (None = store everything significant, 0 = recompute every SCF iteration).
+        Memory budget (GB) for the stored three-center integrals: the near-field blocks with
+        DF_algo=12, every significant block with DF_algo=11 (None = store everything significant,
+        0 = recompute every SCF iteration). The far-field moments of DF_algo=12 are always stored;
+        ``multipole_options['low_memory']`` reduces them on the CPU.
 
     XC_algo : int
         Algorithm selector for XC evaluation (2 for CPU, 3 for GPU).
@@ -313,24 +317,28 @@ class DFT:
         """ Use rys quadrature for the evaluation of two electron integrals (with and without DF)
         In case of DF, only rys quadrature based evaluation of ERIs is supported."""
 
-        self.DF_algo = 11
+        self.DF_algo = 12
         """ This is only for developers. Users should not change it.
-        DF_algo=11 (default, CPU/GPU): the significant 3c2e integrals are evaluated shell-blocked, stored
-        block-sparse in memory throughout the SCF (or partially/never, see max_memory_ints3c2e).
-        DF_algo=10: the previous default (per-function evaluation, sparse triangular storage); it is
+        DF_algo=12 (default, CPU/GPU, pure functionals): the near-field three-center integrals are
+        evaluated shell-blocked and stored as in DF_algo=11, while well-separated (distribution,
+        auxiliary function) pairs are treated through multipole expansions (exact finite moments of
+        the primitive pair products and of the auxiliary functions, box-level expansions truncated
+        at ``lmax``); see `multipole_options`, pyfock.Integrals.df_algo12_helpers and, for the CUDA
+        driver, pyfock.Integrals.df_algo12_helpers_cupy. The 'low_memory' option is CPU only.
+        DF_algo=11 (CPU/GPU): the same shell-blocked evaluation without the far field, i.e. every
+        significant 3c2e integral is stored block-sparse in memory throughout the SCF (or
+        partially/never, see max_memory_ints3c2e). RI exact exchange needs those complete blocks,
+        so a run with xc='HF' or a hybrid functional falls back from 12 to 11 automatically.
+        DF_algo=10: the former default (per-function evaluation, sparse triangular storage); it is
         selectable on both CPU and GPU. Alternatives 1 and 2 are only for reference and take up
         a lot of memory as the complete 3c2e tensor is stored in memory.
         RI-HF (xc='HF') works with DF_algo=1, 2, 3 and 11 (CPU). With 11 the screened blocks are
         orthonormalized in the fit metric once after the integral build (true spherical fit space
-        in SAO mode) and both J and the exchange matrix K are contracted from these rows.
-        DF_algo=12 (CPU, pure functionals): the near-field three-center integrals are evaluated and
-        stored as in DF_algo=11, while well-separated (distribution, auxiliary function) pairs are
-        treated through multipole expansions (exact finite moments of the primitive pair products and
-        of the auxiliary functions, box-level expansions truncated at ``lmax``); see
-        `multipole_options` and pyfock.Integrals.df_algo12_helpers. On the GPU it falls back to 11."""
+        in SAO mode) and both J and the exchange matrix K are contracted from these rows."""
 
         self.max_memory_ints3c2e = None
-        """ Memory budget in GB for the screened three-center integrals when DF_algo=11.
+        """ Memory budget in GB for the screened three-center integrals when DF_algo=11 or 12
+        (with 12 it applies to the near-field blocks; the far-field moments are always stored).
         None (default): keep every significant shell-pair block in memory. A smaller budget keeps the
         most expensive blocks and recomputes the remaining ones in every SCF iteration; 0 recomputes
         everything (direct DF-J). On the GPU this caps cached device values; a separate bounded
@@ -1208,10 +1216,6 @@ class DFT:
 
         if isDF==False:
             strict_schwarz = False
-        if isDF and DF_algo == 12 and self.use_gpu:
-            print('DF_algo=12 (multipole-accelerated density fitting) is implemented for the CPU only; using DF_algo=11 on the GPU.', flush=True)
-            DF_algo = 11
-
         # Skala, the neural functional, is not a LibXC functional: it resolves to a loaded TorchScript
         # model rather than to a list of functional IDs, and is evaluated by its own driver further down.
         # `xc` therefore stays the name string and the LibXC resolution below is skipped.
@@ -1304,6 +1308,13 @@ class DFT:
             else:
                 exx_coef = sum(XC.get_exx_coefficient(fid) for fid in xc)
         self.exx_coef = exx_coef
+        if exx_coef > 0 and isDF and DF_algo == 12:
+            # The far field of DF_algo=12 replaces whole (ij|P) blocks by multipole expansions,
+            # while RI exact exchange contracts the blocks themselves. Fall back to the near-field
+            # algorithm, which is DF_algo=12 without the far field.
+            print('RI exact exchange (HF and hybrid functionals) needs the complete three-center '
+                  'blocks: using DF_algo=11 instead of the default DF_algo=12.', flush=True)
+            DF_algo = 11
         if exx_coef > 0 and isDF and DF_algo not in (1, 2, 3, 11):
             print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
             print('ERROR: RI exact exchange (HF and hybrid functionals) is currently implemented only for DF_algo=1, 2, 3, or 11!')
@@ -1408,11 +1419,11 @@ class DFT:
         
         if strict_schwarz:
             if not (DF_algo in (6, 10, 11, 12)):
-                print('Warning: The stricter variation of Schwarz screening is only compatible with DF algo #6 or #10 so turning it off.')
+                print('Warning: The stricter variation of Schwarz screening is only compatible with DF algo #6, #10, #11 or #12 so turning it off.')
                 strict_schwarz = False
         if cholesky:
             if not (DF_algo in (6, 10, 11, 12)):
-                print('Warning: The Cholesky decomposition of 2c2e integrals is only compatible with DF algo #6 or #10 so turning it off.')
+                print('Warning: The Cholesky decomposition of 2c2e integrals is only compatible with DF algo #6, #10, #11 or #12 so turning it off.')
                 cholesky = False
         if cholesky:
             if self.use_gpu:
@@ -1455,7 +1466,7 @@ class DFT:
         # DF_algo = 8 (no longer works or maintained) # Similar to 6, except that here the significant indices are not stored resulting in 50% memory savings. The drawback is that it only works in serial which is useful for Google colab or Kaggle perhaps.
         # DF_algo = 9 (no longer works or maintained) # 
         # DF_algo = 11 # Shell-blocked Rys evaluation with block-sparse storage and a memory budget (CPU/GPU).
-        # DF_algo = 12 # DF_algo=11 for the near field plus multipole expansions for the far-field three-center integrals (CPU).
+        # DF_algo = 12 # DF_algo=11 for the near field plus multipole expansions for the far-field three-center integrals (CPU/GPU).
         # DF_algo = 10 # Previous default: Similar to 8, but parallelized with the use of Cholesky decomposition for the 2c2e integrals which results in further memory savings and speed up.
 
         V_ecp = None

{pyfock-0.2.0 → pyfock-0.2.1}/pyfock/DFT_Helper_Coulomb.py

@@ -594,32 +594,46 @@ def density_fitting_prelims_for_DFT_development(mol, basis, auxbasis, dftObj, T,
 
         elif DF_algo==12:
             # Near field: shell-blocked Rys blocks (as DF_algo=11); far field: multipole expansions
-            # (Integrals.df_algo12_helpers).  CPU only.
-            if use_gpu:
-                raise NotImplementedError('DF_algo=12 is implemented for the CPU only (DFT.scf falls back to DF_algo=11 on the GPU).')
+            # (Integrals.df_algo12_helpers; Integrals.df_algo12_helpers_cupy on the GPU).
             if rihf:
                 raise NotImplementedError('RI exact exchange needs the complete three-center blocks: use DF_algo=11 for HF and hybrid functionals.')
             print('\n\nPerforming Schwarz screening (multipole-accelerated DF_algo=12)...')
             print('Threshold ', threshold_schwarz)
             startSchwarz = timer()
             start_4c2e_diag = timer()
-            ints4c2e_diag = Integrals.schwarz_helpers.eri_4c2e_diag(basis)
-            sqrt_ints4c2e_diag = np.sqrt(np.abs(ints4c2e_diag))
-            sqrt_diag_ints2c2e = np.sqrt(np.abs(np.diag(ints2c2e)))
+            if use_gpu:
+                if max(basis.shells) - 1 > 4:
+                    # As in DF_algo=11: the legacy CUDA Schwarz diagonal silently omits pairs
+                    # requiring >10 roots; the CPU diagonal has an analytic fallback for them.
+                    ints4c2e_diag = cp.asarray(Integrals.schwarz_helpers.eri_4c2e_diag(basis))
+                else:
+                    ints4c2e_diag = Integrals.df_algo10_helpers_cupy.eri_4c2e_diag_cupy(basis, cp_stream=cp.cuda.get_current_stream())
+                xp = cp
+                build_plan = Integrals.df_algo12_helpers_cupy.build_plan_cupy
+                if not keep_ints3c2e_in_gpu:
+                    print('DF_algo=12 keeps the cached near-field blocks and the far-field moments on the GPU; '
+                          'use max_memory_ints3c2e to limit device storage.', flush=True)
+            else:
+                ints4c2e_diag = Integrals.schwarz_helpers.eri_4c2e_diag(basis)
+                xp = np
+                build_plan = Integrals.df_algo12_helpers.build_plan
+            sqrt_ints4c2e_diag = xp.sqrt(xp.abs(ints4c2e_diag))
+            sqrt_diag_ints2c2e = xp.sqrt(xp.abs(xp.diag(ints2c2e)))
             print('Time taken to evaluate the "diagonal" of 4c2e ERI tensor: ', round(timer() - start_4c2e_diag, 2))
             durationSchwarz = timer() - startSchwarz
             print('Total time taken for Schwarz screening '+str(round(durationSchwarz, 2))+' seconds.\n', flush=True)
             start_plan = timer()
-            ints3c2e = Integrals.df_algo12_helpers.build_plan(basis, auxbasis, sqrt_ints4c2e_diag, sqrt_diag_ints2c2e,
-                                                              threshold_schwarz, strict_schwarz, sao=dftObj.sao,
-                                                              max_memory_gb=dftObj.max_memory_ints3c2e,
-                                                              options=getattr(dftObj, 'multipole_options', None))
+            ints3c2e = build_plan(basis, auxbasis, sqrt_ints4c2e_diag, sqrt_diag_ints2c2e,
+                                  threshold_schwarz, strict_schwarz, sao=dftObj.sao,
+                                  max_memory_gb=dftObj.max_memory_ints3c2e,
+                                  options=getattr(dftObj, 'multipole_options', None))
             print(ints3c2e.summary(), flush=True)
             print('Time taken for the near-field three-center integrals and far-field moments: ', round(timer() - start_plan, 2),
                   ' (' + ', '.join('%s %.2f' % (k, v) for k, v in ints3c2e.timings.items() if k != 'total') + ')', flush=True)
             if strict_schwarz:
                 start_strict_schwarz_nuc_mat = timer()
-                V = Integrals.nuc_mat_symm(basis, mol, None, sqrt_ints4c2e_diag)
+                V = (Integrals.nuc_mat_symm_cupy(basis, mol, None, None, sqrt_ints4c2e_diag)
+                     if use_gpu else Integrals.nuc_mat_symm(basis, mol, None, sqrt_ints4c2e_diag))
                 H = T + V
                 print('Time taken to evaluate the nuclear potential matrix with strict Schwarz screening: ', round(timer() - start_strict_schwarz_nuc_mat, 2), flush=True)
 
@@ -713,6 +727,8 @@ def density_fitting_prelims_for_DFT_development(mol, basis, auxbasis, dftObj, T,
         print('Two Center Two electron ERI size in GB ',ints2c2e.nbytes/1e9, flush=True)
         print('Three Center Two electron ERI (cached near-field blocks) size in GB ', ints3c2e.memory_gb, flush=True)
         print('Far-field multipole moments size in GB ', ints3c2e.moments_gb, flush=True)
+        if use_gpu:
+            print('GPU direct-batch workspace size in GB ', ints3c2e.temporary_memory_gb, flush=True)
     if DF_algo==11:
         print('Two Center Two electron ERI size in GB ',ints2c2e.nbytes/1e9, flush=True)
         if getattr(ints3c2e, 'exchange', None) is not None:
@@ -930,21 +946,26 @@ def Jmat_from_density_fitting(dmat, DF_algo, cholesky, cho_decomp_ints2c2e, df_c
             durationDF_Jtri += timer() - startDF_Jtri
     if DF_algo==12:
         # Near field from the stored blocks, far field through the multipole expansions
-        # (Integrals.df_algo12_helpers); CPU only.
+        # (Integrals.df_algo12_helpers; Integrals.df_algo12_helpers_cupy on the GPU).
         plan = ints3c2e
         startDF_gamma = timer()
-        gamma_alpha = Integrals.df_algo12_helpers.gamma_from_plan(plan, dmat)
+        gamma_alpha = (Integrals.df_algo12_helpers_cupy.gamma_from_plan_cupy(plan, dmat)
+                       if use_gpu else Integrals.df_algo12_helpers.gamma_from_plan(plan, dmat))
         durationDF_gamma += timer() - startDF_gamma
         startDF_coeff = timer()
         with threadpool_limits(limits=ncores, user_api='blas'):
-            if not cholesky:
+            if use_gpu:
+                df_coeff = cp.linalg.solve(ints2c2e, gamma_alpha)
+                cp.cuda.get_current_stream().synchronize()
+            elif not cholesky:
                 df_coeff = scipy.linalg.solve(ints2c2e, gamma_alpha, overwrite_a=False, overwrite_b=False)
             else:
                 df_coeff = scipy.linalg.cho_solve(cho_decomp_ints2c2e, gamma_alpha, overwrite_b=False, check_finite=True)
         durationDF_coeff += timer() - startDF_coeff
-        Ecoul_temp = np.dot(df_coeff, gamma_alpha)
+        Ecoul_temp = (cp.dot(df_coeff, gamma_alpha) if use_gpu else np.dot(df_coeff, gamma_alpha))
         startDF_Jtri = timer()
-        J = Integrals.df_algo12_helpers.J_from_plan(plan, df_coeff)
+        J = (Integrals.df_algo12_helpers_cupy.J_from_plan_cupy(plan, df_coeff)
+             if use_gpu else Integrals.df_algo12_helpers.J_from_plan(plan, df_coeff))
         durationDF_Jtri += timer() - startDF_Jtri
     durationDF = durationDF + timer() - startDF
 

{pyfock-0.2.0 → pyfock-0.2.1}/pyfock/Integrals/__init__.py

@@ -94,6 +94,8 @@ from . import df_algo11_helpers_cupy
 from . import df_algo11_exchange
 from . import multipole_helpers
 from . import df_algo12_helpers
+from . import multipole_helpers_cuda
+from . import df_algo12_helpers_cupy
 from . import integral_helpers
 from .overlap_mat_grad_symm import overlap_mat_grad_symm
 from .kin_mat_grad_symm import kin_mat_grad_symm
@@ -119,7 +121,7 @@ __all__ = ['integral_helpers', 'mmd_nuc_mat_symm', 'nuc_mat_symm', 'ecp_mat_symm
      'rys_4c2e_symm_old', 'rys_4c2e_schwarz_symm', 'conv_3c2e_symm', 'rys_3c2e_symm', 'rys_3c2e_symm_old', 'rys_3c2e_symm_test2', 'conv_2c2e_symm', 'rys_2c2e_symm', 'rys_3c2e_tri', 'rys_nuc_mat_symm', 'schwarz_helpers', 'rys_3c2e_tri_schwarz'\
         'bf_val_helpers', 'eval_xc_1', 'eval_xc_2', 'eval_xc_skala', 'eval_xc_1_cupy', 'eval_xc_2_cupy', 'eval_xc_3_cupy', 'dipole_moment_mat_symm', 'kin_mat_symm_cupy'\
         , 'overlap_mat_symm_cupy', 'dipole_moment_mat_symm_cupy', 'nuc_mat_symm_cupy', 'kin_mat_symm_shell_cupy', 'rys_2c2e_symm_cupy', 'rys_3c2e_symm_cupy',\
-        'rys_3c2e_symm_cupy_fp32', 'schwarz_helpers_cupy', 'df_algo10_helpers', 'df_algo10_helpers_cupy', 'df_algo11_helpers', 'df_algo11_helpers_cupy', 'df_algo11_exchange', 'multipole_helpers', 'df_algo12_helpers', 'overlap_mat_grad_symm', 'kin_mat_grad_symm', 'overlap_mat_grad_r_symm', 'kin_mat_grad_r_symm', 'nuc_mat_grad_symm', 'nuc_mat_grad_r_symm', 'cross_overlap_mat_symm', 'rys_coulomb_matrix',\
+        'rys_3c2e_symm_cupy_fp32', 'schwarz_helpers_cupy', 'df_algo10_helpers', 'df_algo10_helpers_cupy', 'df_algo11_helpers', 'df_algo11_helpers_cupy', 'df_algo11_exchange', 'multipole_helpers', 'df_algo12_helpers', 'multipole_helpers_cuda', 'df_algo12_helpers_cupy', 'overlap_mat_grad_symm', 'kin_mat_grad_symm', 'overlap_mat_grad_r_symm', 'kin_mat_grad_r_symm', 'nuc_mat_grad_symm', 'nuc_mat_grad_r_symm', 'cross_overlap_mat_symm', 'rys_coulomb_matrix',\
          'rys_4c2e_schwarz_sparse_symm', 'rys_coulomb_matrix_sparse', 'rys_2c2e_diag', 'os_4c2e_symm', 'os_4c2e_symm_old', 'os_3c2e_symm', 'rys_3c2e_grad_symm',\
          'rys_3c2e_grad_contract', 'rys_2c2e_grad_contract', 'rys_nuc_grad_contract', 'eval_xc_grad_2', 'eval_xc_grad_skala',\
          'overlap_mat_grad_r_symm_cupy', 'kin_mat_grad_r_symm_cupy', 'nuc_mat_grad_r_symm_cupy', 'ecp_grad_contract']

{pyfock-0.2.0 → pyfock-0.2.1}/pyfock/Integrals/bf_val_helpers.py

@@ -14,6 +14,7 @@ except Exception as e:
         return decorator
 from numba import cuda
 import math
+from .cuda_compat import imax, imax3, imin
 
 def pack_bfs_data(basis):
     """Pack the per-basis-function data of ``basis`` into the NumPy arrays the
@@ -1102,7 +1103,7 @@ def nonzero_ao_indices_batch_cuda(coords, bfs_coords, bfs_radius_cutoff, nblocks
         coord_bf = bfs_coords[ibf]
         offset = iblock*blocksize
         # Loop over the grid points and check if the value of the basis function is greater than the threshold
-        for igrd in range(offset,  min(offset+blocksize, ngrids)):
+        for igrd in range(offset,  imin(offset+blocksize, ngrids)):
             coord_grid = coords[igrd]
             x = coord_grid[0]-coord_bf[0]
             y = coord_grid[1]-coord_bf[1]

pyfock-0.2.1/pyfock/Integrals/cuda_compat.py

@@ -0,0 +1,31 @@
+"""Small device helpers that replace Python builtins Numba's CUDA target cannot lower.
+
+Numba implements ``max``/``min`` as a variadic ``@overload`` whose implementation is a
+``def impl(*x)`` device function.  The CUDA target compiles device functions by signature and
+rejects the ``*args`` form ("Signature mismatch: N argument types given, but function takes 1
+arguments"), so *any* use of ``max`` or ``min`` inside a ``@cuda.jit`` function fails to
+compile, whatever the number of arguments.  These helpers are the explicit equivalents; they
+inline to a single comparison and are exact for the integers the kernels use them on.
+"""
+from numba import cuda
+
+__all__ = ['imax', 'imin', 'imax3']
+
+
+@cuda.jit(device=True, cache=True)
+def imax(a, b):
+    """``max(a, b)``."""
+    return a if a > b else b
+
+
+@cuda.jit(device=True, cache=True)
+def imin(a, b):
+    """``min(a, b)``."""
+    return a if a < b else b
+
+
+@cuda.jit(device=True, cache=True)
+def imax3(a, b, c):
+    """``max(a, b, c)``."""
+    m = a if a > b else b
+    return m if m > c else c

{pyfock-0.2.0 → pyfock-0.2.1}/pyfock/Integrals/df_algo10_helpers_cupy.py

@@ -40,6 +40,7 @@ from .rys_helpers_cuda import coulomb_rys, coulomb_rys_3c2e, Roots, Roots_5, DAT
 from .df_algo10_helpers import (
     STRICT_PAIR_CUTOFF, pack_basis_arrays, aux_shell_arrays, sao_aux_projectors,
     check_sao_bounds_are_shell_constant)
+from .cuda_compat import imax, imax3, imin
 
 __all__ = [
     'eri_4c2e_diag_cupy',
@@ -129,8 +130,8 @@ def rys_eri_4c2e_diag_internal_cuda(bfs_coords, bfs_contr_prim_norms, bfs_lmn, b
         ld, md, nd = lmnj
 
         norder = int((la + ma + na + lb + mb + nb + lc + mc + nc + ld + md + nd) / 2 + 1)
-        n = int(max(la + lb, ma + mb, na + nb))
-        m = int(max(lc + ld, mc + md, nc + nd))
+        n = int(imax3(la + lb, ma + mb, na + nb))
+        m = int(imax3(lc + ld, mc + md, nc + nd))
         roots = cuda.local.array((10), numba.float64)
         weights = cuda.local.array((10), numba.float64)
         G = cuda.local.array((13, 13), numba.float64)
@@ -330,8 +331,8 @@ def rys_3c2e_tri_schwarz_sparse_algo10_internal_cuda(bfs_coords, bfs_contr_prim_
                     val = 0.0
 
                     if norder <= 10:
-                        n = int(max(la + lb, ma + mb, na + nb))
-                        m = int(max(lc + ld, mc + md, nc + nd))
+                        n = int(imax3(la + lb, ma + mb, na + nb))
+                        m = int(imax3(lc + ld, mc + md, nc + nd))
 
                         for kk in range(nprimk):
                             dkk = aux_bfs_coeffs[k][kk]
@@ -485,8 +486,8 @@ def rys_3c2e_tri_schwarz_sparse_algo10_sao_internal_cuda(bfs_coords, bfs_contr_p
                     val = 0.0
 
                     if norder <= 10:
-                        n = int(max(la + lb, ma + mb, na + nb))
-                        m = int(max(lc + ld, mc + md, nc + nd))
+                        n = int(imax3(la + lb, ma + mb, na + nb))
+                        m = int(imax3(lc + ld, mc + md, nc + nd))
 
                         for kk in range(nprimk):
                             dkk = aux_bfs_coeffs[k][kk]

{pyfock-0.2.0 → pyfock-0.2.1}/pyfock/Integrals/df_algo11_cuda_core.py

@@ -4,6 +4,21 @@ Small items use one thread; large items use one block, shared Rys/shift tables,
 and strided component accumulators. No integral is evaluated per AO function.
 The build writes Cartesian values; a separate row-owned pass projects whole
 auxiliary shells, including h/i shells, without in-place read/write races.
+
+``evaluate_item`` serves both density-fitting algorithms. With ``masked`` false
+(algorithm 11) every surviving primitive pair of the shell pair contributes and
+``mask`` is ignored; with ``masked`` true (algorithm 12) a primitive pair is
+skipped for auxiliary shell ``K`` when its branch is far field for that shell,
+so the block holds the near-field part only and the rest comes from the
+multipole expansions. ``mask`` is ``(pp_off, pp_branch, ff)``: the branch of
+primitive pair ``(ip, jp)`` of pair ``p`` is ``pp_branch[pp_off[p] + ip * nprimB
++ jp]`` and ``ff[branch, K]`` is its far-field flag. The index is dense in
+``(ip, jp)`` rather than a count of surviving pairs on purpose: the device would
+otherwise have to reproduce the host's Gaussian-product pre-screen decision
+exactly, and a one-ulp disagreement would shift every later branch of that shell
+pair. Pairs the host dropped carry a branch whose row is far field everywhere, so
+they are skipped here too even if this pre-screen keeps them. Both callers pass
+arrays of the same dtypes, so one specialization is compiled.
 """
 import math
 
@@ -15,9 +30,10 @@ from .rys_helpers_cuda import Roots, Recur_3c2e, Shift_3c2e
 
 @cuda.jit(device=True, cache=True)
 def evaluate_item(item, lane, stride, cooperative, orbital, auxiliary, shells,
-                  aux_shells, pairs, items, offsets, values, data_x, data_w,
+                  aux_shells, pairs, items, offsets, values, data_x, data_w, masked, mask,
                   roots, weights, gx, gy, gz, sx, sy, sz, accum, tmp):
     coords, lmn, nprim, expnts, coef = orbital
+    pp_off, pp_branch, ff = mask
     acoords, almn, anprim, aexpnts, acoef = auxiliary
     off, nbf, ls = shells
     aoff, anbf, als = aux_shells
@@ -43,6 +59,8 @@ def evaluate_item(item, lane, stride, cooperative, orbital, auxiliary, shells,
             ab = alpha * beta
             if ab / gp * r2 > EXP_ARG_CUTOFF:
                 continue
+            if masked and ff[pp_branch[pp_off[p] + ip * nprim[b0] + jp], K]:
+                continue
             px = (alpha * coords[a0, 0] + beta * coords[b0, 0]) / gp
             py = (alpha * coords[a0, 1] + beta * coords[b0, 1]) / gp
             pz = (alpha * coords[a0, 2] + beta * coords[b0, 2]) / gp