Skip to content

Commit 8bbf73b

Browse files
marco-2023marco-2023
authored
Merge pull request #211 from ohzde/screening
one index screening
2 parents 9f4b422 + 05fcf65 commit 8bbf73b

File tree

7 files changed

+330
-69
lines changed

7 files changed

+330
-69
lines changed

‎gbasis/evals/density.py‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,15 @@ def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval):
6262
return np.sum(density, axis=0)
6363

6464

65-
def evaluate_density(one_density_matrix, basis, points, transform=None, threshold=1.0e-8):
65+
def evaluate_density(
66+
one_density_matrix,
67+
basis,
68+
points,
69+
transform=None,
70+
threshold=1.0e-8,
71+
screen_basis=True,
72+
tol_screen=1e-8,
73+
):
6674
r"""Return the density of the given basis set at the given points.
6775
6876
.. math::
@@ -94,14 +102,23 @@ def evaluate_density(one_density_matrix, basis, points, transform=None, threshol
94102
threshold : float, optional
95103
The absolute value below which negative density values are acceptable. Any negative density
96104
value with an absolute value smaller than this threshold will be set to zero.
105+
screen_basis : bool, optional
106+
A toggle to enable or disable screening. Default value is True (enable screening).
107+
tol_screen : float, optional
108+
The tolerance used for screening overlap integrals. `tol_screen` is combined with the
109+
minimum contraction exponents to compute a cutoff radius which is compared to the distance
110+
between the points and the contraction centers to decide whether the basis function
111+
should be evaluated or set to zero at that point.
97112
98113
Returns
99114
-------
100115
density : np.ndarray(N,)
101116
Density evaluated at `N` grid points.
102117
103118
"""
104-
orb_eval = evaluate_basis(basis, points, transform=transform)
119+
orb_eval = evaluate_basis(
120+
basis, points, transform=transform, screen_basis=screen_basis, tol_screen=tol_screen
121+
)
105122
output = evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval)
106123
# Fix #117: check magnitude of small negative density values, then use clip to remove them
107124
min_output = np.min(output)
@@ -489,7 +506,7 @@ def evaluate_density_hessian(
489506
r"""Return the Hessian of the density evaluated at the given points.
490507
491508
.. math::
492-
509+
493510
H[\rho(\mathbf{r})]
494511
=
495512
\begin{bmatrix}

‎gbasis/evals/eval.py‎

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from gbasis.base_one import BaseOneIndex
33
from gbasis.contractions import GeneralizedContractionShell
44
from gbasis.evals._deriv import _eval_deriv_contractions
5+
from gbasis.screening import get_points_mask_for_contraction
56
import numpy as np
67

78

@@ -53,7 +54,7 @@ class Eval(BaseOneIndex):
5354
"""
5455

5556
@staticmethod
56-
def construct_array_contraction(contractions, points):
57+
def construct_array_contraction(contractions, points, screen_basis=True, tol_screen=1e-8):
5758
r"""Return the evaluations of the given contractions at the given coordinates.
5859
5960
Parameters
@@ -105,13 +106,29 @@ def construct_array_contraction(contractions, points):
105106
angmom_comps = contractions.angmom_components_cart
106107
center = contractions.coord
107108
norm_prim_cart = contractions.norm_prim_cart
108-
output = _eval_deriv_contractions(
109-
points, np.zeros(3), center, angmom_comps, alphas, prim_coeffs, norm_prim_cart
109+
110+
# if screening is not requested, evaluate all points
111+
if not screen_basis:
112+
return _eval_deriv_contractions(
113+
points, np.zeros(3), center, angmom_comps, alphas, prim_coeffs, norm_prim_cart
114+
)
115+
116+
# default case, screen points that are too far from the contraction center
117+
points_mask = get_points_mask_for_contraction(contractions, points, tol_screen=tol_screen)
118+
# reconstruct the array with correct shape
119+
L = angmom_comps.shape[0]
120+
M = prim_coeffs.shape[1]
121+
N = points.shape[0]
122+
output = np.zeros((M, L, N), dtype=np.float64)
123+
124+
# fill non-screened points in the output array
125+
output[:,:, points_mask] = _eval_deriv_contractions(
126+
points[points_mask], np.zeros(3), center, angmom_comps, alphas, prim_coeffs, norm_prim_cart
110127
)
111128
return output
112129

113130

114-
def evaluate_basis(basis, points, transform=None):
131+
def evaluate_basis(basis, points, transform=None, screen_basis=True, tol_screen=1e-8):
115132
r"""Evaluate the basis set in the given coordinate system at the given points.
116133
117134
Parameters
@@ -129,6 +146,13 @@ def evaluate_basis(basis, points, transform=None):
129146
Transformation is applied to the left, i.e. the sum is over the index 1 of `transform`
130147
and index 0 of the array for contractions.
131148
Default is no transformation.
149+
screen_basis : bool, optional
150+
A toggle to enable or disable screening. Default value is True (enable screening).
151+
tol_screen : float, optional
152+
The tolerance used for screening overlap integrals. `tol_screen` is combined with the
153+
minimum contraction exponents to compute a cutoff radius which is compared to the distance
154+
between the points and the contraction centers to decide whether the basis function
155+
should be evaluated or set to zero at that point.
132156
133157
Returns
134158
-------
@@ -141,11 +165,12 @@ def evaluate_basis(basis, points, transform=None):
141165
142166
"""
143167
coord_type = [ct for ct in [shell.coord_type for shell in basis]]
168+
kwargs = {"tol_screen": tol_screen, "screen_basis": screen_basis}
144169

145170
if transform is not None:
146-
return Eval(basis).construct_array_lincomb(transform, coord_type, points=points)
171+
return Eval(basis).construct_array_lincomb(transform, coord_type, points=points, **kwargs)
147172
if all(ct == "cartesian" for ct in coord_type):
148-
return Eval(basis).construct_array_cartesian(points=points)
173+
return Eval(basis).construct_array_cartesian(points=points, **kwargs)
149174
if all(ct == "spherical" for ct in coord_type):
150-
return Eval(basis).construct_array_spherical(points=points)
151-
return Eval(basis).construct_array_mix(coord_type, points=points)
175+
return Eval(basis).construct_array_spherical(points=points, **kwargs)
176+
return Eval(basis).construct_array_mix(coord_type, points=points, **kwargs)

‎gbasis/screening.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
"""1 and 2-index screening functions"""
22

33
import numpy as np
4+
from scipy.special import lambertw
5+
from gbasis.utils import factorial2
46

57

68
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
4244
cutoff = np.sqrt(-(alpha_a + alpha_b) / (alpha_a * alpha_b) * np.log(tol_screen))
4345
# integrals are screened if centers are further apart than the cutoff
4446
return np.linalg.norm(r_12) > cutoff
47+
48+
49+
def get_points_mask_for_contraction(contractions, points, tol_screen):
50+
r"""Return a boolean mask indicating which points should be screened for a contraction shell
51+
52+
A point is considered screened if it lies farther from the contraction center than a cutoff
53+
radius computed from the contraction parameters and the screening tolerance
54+
(see :func:`compute_primitive_cutoff_radius`).
55+
56+
Parameters
57+
----------
58+
contractions : GeneralizedContractionShell
59+
Contracted Cartesian Gaussians (of the same shell) for which the mask is computed.
60+
points : np.ndarray of shape (N, 3)
61+
Cartesian coordinates of the points in space (in atomic units) where the
62+
basis functions are evaluated. Rows correspond to points; columns
63+
correspond to the :math:`x`, :math:`y`, and :math:`z` components.
64+
tol_screen : float
65+
Screening tolerance for excluding points. This value, together with the
66+
minimum contraction parameters, determines a cutoff distance. Points
67+
farther than the cutoff are excluded from contraction evaluation.
68+
69+
Returns
70+
-------
71+
np.ndarray of bool, shape (N,)
72+
Boolean mask where `False` marks points to be screened out.
73+
"""
74+
angm = contractions.angmom
75+
# reshape exponents for broadcasting
76+
exps = contractions.exps[:, np.newaxis] # shape (K, 1)
77+
# use absolute value (indicating magnitude) of primitive contributions
78+
coeffs = np.abs(contractions.coeffs) # shape (K, M)
79+
80+
# compute cutoff radius for all primitives in all contractions
81+
r_cuts = compute_primitive_cutoff_radius(coeffs, exps, angm, tol_screen)
82+
83+
# pick the maximum radius over all primitives and contractions
84+
cutoff_radius = np.max(r_cuts)
85+
86+
# screen out points beyond the cutoff radius
87+
points_r = np.linalg.norm(points - contractions.coord, axis=1)
88+
return points_r <= cutoff_radius
89+
90+
91+
def compute_primitive_cutoff_radius(c, alpha, angm, tol_screen):
92+
r"""Compute the cutoff radius for a primitive Gaussian.
93+
94+
The cutoff radius is the maximum distance from the center of the primitive Gaussian at which the
95+
radial density remains above a given tolerance :math:`\epsilon`. This radius is computed by
96+
solving the equation:
97+
98+
.. math::
99+
100+
r^{2\ell} \chi(r)^2 = \epsilon^2
101+
102+
where :math:`\chi(r)` is the radial part of the primitive Gaussian, defined as:
103+
104+
.. math::
105+
106+
\chi(r) = c \, n \, e^{-\alpha r^2}
107+
108+
Here, :math:`c` is the coefficient of the primitive Gaussian, :math:`\alpha` is its exponent,
109+
:math:`\ell` is the angular momentum quantum number, and :math:`n` is the normalization factor
110+
given by:
111+
112+
.. math::
113+
114+
n = \left( \frac{2 \alpha}{\pi} \right)^{\frac{1}{4}}
115+
\frac{(4 \alpha)^{\frac{\ell}{2}}}{\sqrt{(2\ell + 1)!!}}
116+
117+
Parameters
118+
----------
119+
c : float
120+
Coefficient of the primitive Gaussian.
121+
alpha : float
122+
Exponent :math:`\alpha` of the primitive Gaussian.
123+
angm : int
124+
Angular momentum quantum number :math:`\ell` of the primitive Gaussian.
125+
dens_tol : float
126+
Radial density tolerance :math:`\epsilon` for computing the cutoff radius.
127+
128+
Returns
129+
-------
130+
float
131+
The cutoff radius for the primitive Gaussian.
132+
"""
133+
# Compute normalization factor n for the primitive Gaussian
134+
n = (2 * alpha / np.pi) ** 0.25 * (4 * alpha) ** (angm / 2) / np.sqrt(factorial2(2 * angm + 1))
135+
# special case for angular momentum 0. Solution found using logarithm
136+
if angm == 0:
137+
return np.sqrt(-np.log(tol_screen / (c * n)) / alpha)
138+
# general case for angular momentum > 0. Solution found in terms of the Lambert W function
139+
# W_{-1} branch corresponds to the outermost solution
140+
lambert_input_value = -2 * alpha * (tol_screen / (c * n)) ** (2 / angm) / angm
141+
return np.sqrt(-(angm / (2 * alpha)) * lambertw(lambert_input_value, k=-1).real)

‎tests/test_density.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ def test_evaluate_density():
6868
density += density.T
6969
points = np.random.rand(10, 3)
7070

71-
evaluate_orbs = evaluate_basis(basis, points, transform)
71+
evaluate_orbs = evaluate_basis(basis, points, transform, screen_basis=False)
7272
dens = evaluate_density(density, basis, points, transform)
7373
assert np.all(dens >= 0.0)
7474
assert np.allclose(dens, np.einsum("ij,ik,jk->k", density, evaluate_orbs, evaluate_orbs))
@@ -89,12 +89,12 @@ def test_evaluate_deriv_density():
8989
"ij,ik,jk->k",
9090
density,
9191
evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
92-
evaluate_basis(basis, points, transform),
92+
evaluate_basis(basis, points, transform, screen_basis=False),
9393
)
9494
+ np.einsum(
9595
"ij,ik,jk->k",
9696
density,
97-
evaluate_basis(basis, points, transform),
97+
evaluate_basis(basis, points, transform, screen_basis=False),
9898
evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
9999
),
100100
)
@@ -105,12 +105,12 @@ def test_evaluate_deriv_density():
105105
"ij,ik,jk->k",
106106
density,
107107
evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
108-
evaluate_basis(basis, points, transform),
108+
evaluate_basis(basis, points, transform, screen_basis=False),
109109
)
110110
+ np.einsum(
111111
"ij,ik,jk->k",
112112
density,
113-
evaluate_basis(basis, points, transform),
113+
evaluate_basis(basis, points, transform, screen_basis=False),
114114
evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
115115
),
116116
)
@@ -121,12 +121,12 @@ def test_evaluate_deriv_density():
121121
"ij,ik,jk->k",
122122
density,
123123
evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
124-
evaluate_basis(basis, points, transform),
124+
evaluate_basis(basis, points, transform, screen_basis=False),
125125
)
126126
+ np.einsum(
127127
"ij,ik,jk->k",
128128
density,
129-
evaluate_basis(basis, points, transform),
129+
evaluate_basis(basis, points, transform, screen_basis=False),
130130
evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
131131
),
132132
)
@@ -235,36 +235,36 @@ def test_evaluate_density_gradient():
235235
"ij,ik,jk->k",
236236
density,
237237
evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
238-
evaluate_basis(basis, points, transform),
238+
evaluate_basis(basis, points, transform, screen_basis=False),
239239
)
240240
+ np.einsum(
241241
"ij,ik,jk->k",
242242
density,
243-
evaluate_basis(basis, points, transform),
243+
evaluate_basis(basis, points, transform, screen_basis=False),
244244
evaluate_deriv_basis(basis, points, np.array([1, 0, 0]), transform),
245245
),
246246
np.einsum(
247247
"ij,ik,jk->k",
248248
density,
249249
evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
250-
evaluate_basis(basis, points, transform),
250+
evaluate_basis(basis, points, transform, screen_basis=False),
251251
)
252252
+ np.einsum(
253253
"ij,ik,jk->k",
254254
density,
255-
evaluate_basis(basis, points, transform),
255+
evaluate_basis(basis, points, transform, screen_basis=False),
256256
evaluate_deriv_basis(basis, points, np.array([0, 1, 0]), transform),
257257
),
258258
np.einsum(
259259
"ij,ik,jk->k",
260260
density,
261261
evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
262-
evaluate_basis(basis, points, transform),
262+
evaluate_basis(basis, points, transform, screen_basis=False),
263263
)
264264
+ np.einsum(
265265
"ij,ik,jk->k",
266266
density,
267-
evaluate_basis(basis, points, transform),
267+
evaluate_basis(basis, points, transform, screen_basis=False),
268268
evaluate_deriv_basis(basis, points, np.array([0, 0, 1]), transform),
269269
),
270270
]

0 commit comments

Comments
 (0)