one index screening

gbasis/evals/density.py

@@ -62,7 +62,15 @@ def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval):
     return np.sum(density, axis=0)
 
 
-def evaluate_density(one_density_matrix, basis, points, transform=None, threshold=1.0e-8):
+def evaluate_density(
+    one_density_matrix,
+    basis,
+    points,
+    transform=None,
+    threshold=1.0e-8,
+    screen_basis=True,
+    tol_screen=1e-8,
+):
     r"""Return the density of the given basis set at the given points.
 
     .. math::
@@ -94,14 +102,23 @@ def evaluate_density(one_density_matrix, basis, points, transform=None, threshol
     threshold : float, optional
         The absolute value below which negative density values are acceptable. Any negative density
         value with an absolute value smaller than this threshold will be set to zero.
+    screen_basis : bool, optional
+        A toggle to enable or disable screening. Default value is True (enable screening).
+    tol_screen : float, optional
+        The tolerance used for screening overlap integrals. `tol_screen` is combined with the
+        minimum contraction exponents to compute a cutoff radius which is compared to the distance
+        between the points and the contraction centers to decide whether the basis function
+        should be evaluated or set to zero at that point.
 
     Returns
     -------
     density : np.ndarray(N,)
         Density evaluated at `N` grid points.
 
     """
-    orb_eval = evaluate_basis(basis, points, transform=transform)
+    orb_eval = evaluate_basis(
+        basis, points, transform=transform, screen_basis=screen_basis, tol_screen=tol_screen
+    )
     output = evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval)
     # Fix #117: check magnitude of small negative density values, then use clip to remove them
     min_output = np.min(output)
@@ -489,7 +506,7 @@ def evaluate_density_hessian(
     r"""Return the Hessian of the density evaluated at the given points.
 
     .. math::
-        
+
         H[\rho(\mathbf{r})]
         =
         \begin{bmatrix}

gbasis/evals/eval.py

@@ -2,6 +2,7 @@
 from gbasis.base_one import BaseOneIndex
 from gbasis.contractions import GeneralizedContractionShell
 from gbasis.evals._deriv import _eval_deriv_contractions
+from gbasis.screening import get_points_mask_for_contraction
 import numpy as np
 
 
@@ -53,7 +54,7 @@ class Eval(BaseOneIndex):
     """
 
     @staticmethod
-    def construct_array_contraction(contractions, points):
+    def construct_array_contraction(contractions, points, screen_basis=True, tol_screen=1e-8):
         r"""Return the evaluations of the given contractions at the given coordinates.
 
         Parameters
@@ -105,13 +106,29 @@ def construct_array_contraction(contractions, points):
         angmom_comps = contractions.angmom_components_cart
         center = contractions.coord
         norm_prim_cart = contractions.norm_prim_cart
-        output = _eval_deriv_contractions(
-            points, np.zeros(3), center, angmom_comps, alphas, prim_coeffs, norm_prim_cart
+
+        # if screening is not requested, evaluate all points
+        if not screen_basis:
+            return _eval_deriv_contractions(
+                points, np.zeros(3), center, angmom_comps, alphas, prim_coeffs, norm_prim_cart
+            )
+
+        # default case, screen points that are too far from the contraction center
+        points_mask = get_points_mask_for_contraction(contractions, points, tol_screen=tol_screen)
+        # reconstruct the array with correct shape
+        L = angmom_comps.shape[0]
+        M = prim_coeffs.shape[1]
+        N = points.shape[0]
+        output = np.zeros((M, L, N), dtype=np.float64)
+
+        # fill non-screened points in the output array
+        output[:,:, points_mask] = _eval_deriv_contractions(
+            points[points_mask], np.zeros(3), center, angmom_comps, alphas, prim_coeffs, norm_prim_cart
         )
         return output
 
 
-def evaluate_basis(basis, points, transform=None):
+def evaluate_basis(basis, points, transform=None, screen_basis=True, tol_screen=1e-8):
     r"""Evaluate the basis set in the given coordinate system at the given points.
 
     Parameters
@@ -129,6 +146,13 @@ def evaluate_basis(basis, points, transform=None):
         Transformation is applied to the left, i.e. the sum is over the index 1 of `transform`
         and index 0 of the array for contractions.
         Default is no transformation.
+    screen_basis : bool, optional
+        A toggle to enable or disable screening. Default value is True (enable screening).
+    tol_screen : float, optional
+        The tolerance used for screening overlap integrals. `tol_screen` is combined with the
+        minimum contraction exponents to compute a cutoff radius which is compared to the distance
+        between the points and the contraction centers to decide whether the basis function
+        should be evaluated or set to zero at that point.
 
     Returns
     -------
@@ -141,11 +165,12 @@ def evaluate_basis(basis, points, transform=None):
 
     """
     coord_type = [ct for ct in [shell.coord_type for shell in basis]]
+    kwargs = {"tol_screen": tol_screen, "screen_basis": screen_basis}
 
     if transform is not None:
-        return Eval(basis).construct_array_lincomb(transform, coord_type, points=points)
+        return Eval(basis).construct_array_lincomb(transform, coord_type, points=points, **kwargs)
     if all(ct == "cartesian" for ct in coord_type):
-        return Eval(basis).construct_array_cartesian(points=points)
+        return Eval(basis).construct_array_cartesian(points=points, **kwargs)
     if all(ct == "spherical" for ct in coord_type):
-        return Eval(basis).construct_array_spherical(points=points)
-    return Eval(basis).construct_array_mix(coord_type, points=points)
+        return Eval(basis).construct_array_spherical(points=points, **kwargs)
+    return Eval(basis).construct_array_mix(coord_type, points=points, **kwargs)

gbasis/screening.py

@@ -1,6 +1,8 @@
 """1 and 2-index screening functions"""
 
 import numpy as np
+from scipy.special import lambertw
+from gbasis.utils import factorial2
 
 
 def is_two_index_overlap_screened(contractions_one, contractions_two, tol_screen):
@@ -42,3 +44,98 @@ def is_two_index_overlap_screened(contractions_one, contractions_two, tol_screen
     cutoff = np.sqrt(-(alpha_a + alpha_b) / (alpha_a * alpha_b) * np.log(tol_screen))
     # integrals are screened if centers are further apart than the cutoff
     return np.linalg.norm(r_12) > cutoff
+
+
+def get_points_mask_for_contraction(contractions, points, tol_screen):
+    r"""Return a boolean mask indicating which points should be screened for a contraction shell
+
+    A point is considered screened if it lies farther from the contraction center than a cutoff
+    radius computed from the contraction parameters and the screening tolerance
+    (see :func:`compute_primitive_cutoff_radius`).
+
+    Parameters
+    ----------
+    contractions : GeneralizedContractionShell
+        Contracted Cartesian Gaussians (of the same shell) for which the mask is computed.
+    points : np.ndarray of shape (N, 3)
+        Cartesian coordinates of the points in space (in atomic units) where the
+        basis functions are evaluated. Rows correspond to points; columns
+        correspond to the :math:`x`, :math:`y`, and :math:`z` components.
+    tol_screen : float
+        Screening tolerance for excluding points. This value, together with the
+        minimum contraction parameters, determines a cutoff distance. Points
+        farther than the cutoff are excluded from contraction evaluation.
+
+    Returns
+    -------
+    np.ndarray of bool, shape (N,)
+        Boolean mask where `False` marks points to be screened out.
+    """
+    angm = contractions.angmom
+    # reshape exponents for broadcasting
+    exps = contractions.exps[:, np.newaxis]  # shape (K, 1)
+    # use absolute value (indicating magnitude) of primitive contributions
+    coeffs = np.abs(contractions.coeffs)  # shape (K, M)
+
+    # compute cutoff radius for all primitives in all contractions
+    r_cuts = compute_primitive_cutoff_radius(coeffs, exps, angm, tol_screen)
+
+    # pick the maximum radius over all primitives and contractions
+    cutoff_radius = np.max(r_cuts)
+
+    # screen out points beyond the cutoff radius
+    points_r = np.linalg.norm(points - contractions.coord, axis=1)
+    return points_r <= cutoff_radius
+
+
+def compute_primitive_cutoff_radius(c, alpha, angm, tol_screen):
+    r"""Compute the cutoff radius for a primitive Gaussian.
+
+    The cutoff radius is the maximum distance from the center of the primitive Gaussian at which the
+    radial density remains above a given tolerance :math:`\epsilon`. This radius is computed by
+    solving the equation:
+
+    .. math::
+
+        r^{2\ell} \chi(r)^2 = \epsilon^2
+
+    where :math:`\chi(r)` is the radial part of the primitive Gaussian, defined as:
+
+    .. math::
+
+        \chi(r) = c \, n \, e^{-\alpha r^2}
+
+    Here, :math:`c` is the coefficient of the primitive Gaussian, :math:`\alpha` is its exponent,
+    :math:`\ell` is the angular momentum quantum number, and :math:`n` is the normalization factor
+    given by:
+
+    .. math::
+
+        n = \left( \frac{2 \alpha}{\pi} \right)^{\frac{1}{4}}
+            \frac{(4 \alpha)^{\frac{\ell}{2}}}{\sqrt{(2\ell + 1)!!}}
+
+    Parameters
+    ----------
+    c : float
+        Coefficient of the primitive Gaussian.
+    alpha : float
+        Exponent :math:`\alpha` of the primitive Gaussian.
+    angm : int
+        Angular momentum quantum number :math:`\ell` of the primitive Gaussian.
+    dens_tol : float
+        Radial density tolerance :math:`\epsilon` for computing the cutoff radius.
+
+    Returns
+    -------
+    float
+        The cutoff radius for the primitive Gaussian.
+    """
+    # Compute normalization factor n for the primitive Gaussian
+    n = (2 * alpha / np.pi) ** 0.25 * (4 * alpha) ** (angm / 2) / np.sqrt(factorial2(2 * angm + 1))
+    # special case for angular momentum 0. Solution found using logarithm
+    if angm == 0:
+        return np.sqrt(-np.log(tol_screen / (c * n)) / alpha)
+    # general case for angular momentum > 0. Solution found in terms of the Lambert W function
+    # W_{-1} branch corresponds to the outermost solution
+    lambert_input_value = -2 * alpha * (tol_screen / (c * n)) ** (2 / angm) / angm
+    return np.sqrt(-(angm / (2 * alpha)) * lambertw(lambert_input_value, k=-1).real)

tests/test_density.py

@@ -68,7 +68,7 @@ def test_evaluate_density():
     density += density.T
     points = np.random.rand(10, 3)
 
-    evaluate_orbs = evaluate_basis(basis, points, transform)
+    evaluate_orbs = evaluate_basis(basis, points, transform, screen_basis=False)
     dens = evaluate_density(density, basis, points, transform)
     assert np.all(dens >= 0.0)
     assert np.allclose(dens, np.einsum("ij,ik,jk->k", density, evaluate_orbs, evaluate_orbs))
@@ -89,12 +89,12 @@ def test_evaluate_deriv_density():
             "ij,ik,jk->k",
             density,
             evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
-            evaluate_basis(basis, points, transform),
+            evaluate_basis(basis, points, transform, screen_basis=False),
         )
         + np.einsum(
             "ij,ik,jk->k",
             density,
-            evaluate_basis(basis, points, transform),
+            evaluate_basis(basis, points, transform, screen_basis=False),
             evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
         ),
     )
@@ -105,12 +105,12 @@ def test_evaluate_deriv_density():
             "ij,ik,jk->k",
             density,
             evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
-            evaluate_basis(basis, points, transform),
+            evaluate_basis(basis, points, transform, screen_basis=False),
         )
         + np.einsum(
             "ij,ik,jk->k",
             density,
-            evaluate_basis(basis, points, transform),
+            evaluate_basis(basis, points, transform, screen_basis=False),
             evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
         ),
     )
@@ -121,12 +121,12 @@ def test_evaluate_deriv_density():
             "ij,ik,jk->k",
             density,
             evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
-            evaluate_basis(basis, points, transform),
+            evaluate_basis(basis, points, transform, screen_basis=False),
         )
         + np.einsum(
             "ij,ik,jk->k",
             density,
-            evaluate_basis(basis, points, transform),
+            evaluate_basis(basis, points, transform, screen_basis=False),
             evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
         ),
     )
@@ -235,36 +235,36 @@ def test_evaluate_density_gradient():
                     "ij,ik,jk->k",
                     density,
                     evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
-                    evaluate_basis(basis, points, transform),
+                    evaluate_basis(basis, points, transform, screen_basis=False),
                 )
                 + np.einsum(
                     "ij,ik,jk->k",
                     density,
-                    evaluate_basis(basis, points, transform),
+                    evaluate_basis(basis, points, transform, screen_basis=False),
                     evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
                 ),
                 np.einsum(
                     "ij,ik,jk->k",
                     density,
                     evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
-                    evaluate_basis(basis, points, transform),
+                    evaluate_basis(basis, points, transform, screen_basis=False),
                 )
                 + np.einsum(
                     "ij,ik,jk->k",
                     density,
-                    evaluate_basis(basis, points, transform),
+                    evaluate_basis(basis, points, transform, screen_basis=False),
                     evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
                 ),
                 np.einsum(
                     "ij,ik,jk->k",
                     density,
                     evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
-                    evaluate_basis(basis, points, transform),
+                    evaluate_basis(basis, points, transform, screen_basis=False),
                 )
                 + np.einsum(
                     "ij,ik,jk->k",
                     density,
-                    evaluate_basis(basis, points, transform),
+                    evaluate_basis(basis, points, transform, screen_basis=False),
                     evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
                 ),
             ]