diff --git a/gbasis/integrals/_schwarz_screening.py b/gbasis/integrals/_schwarz_screening.py new file mode 100644 index 00000000..306769df --- /dev/null +++ b/gbasis/integrals/_schwarz_screening.py @@ -0,0 +1,218 @@ +"""Integral screening utilities for efficient 2-electron integral computation. + +This module implements Schwarz screening and shell-pair screening to skip +negligible integrals, providing speedup for spatially extended systems. + +References: +- Häser, M. & Ahlrichs, R. J. Comput. Chem. 1989, 10, 104. +- Gill, P. M. W.; Johnson, B. G.; Pople, J. A. Int. J. Quantum Chem. 1991, 40, 745. +""" + +import numpy as np + + +def compute_schwarz_bound_shell_pair(boys_func, cont_one, cont_two, compute_integral_func): + """Compute Schwarz bound for a shell pair: sqrt((ab|ab)). + + Parameters + ---------- + boys_func : callable + Boys function for integral evaluation. + cont_one : GeneralizedContractionShell + First contracted shell. + cont_two : GeneralizedContractionShell + Second contracted shell. + compute_integral_func : callable + Function to compute (ab|cd) integrals. + + Returns + ------- + bound : float + Schwarz bound sqrt(max|(ab|ab)|) for this shell pair. + """ + # Compute (ab|ab) integral + integral = compute_integral_func( + boys_func, + cont_one.coord, + cont_one.angmom, + cont_one.angmom_components_cart, + cont_one.exps, + cont_one.coeffs, + cont_two.coord, + cont_two.angmom, + cont_two.angmom_components_cart, + cont_two.exps, + cont_two.coeffs, + cont_one.coord, + cont_one.angmom, + cont_one.angmom_components_cart, + cont_one.exps, + cont_one.coeffs, + cont_two.coord, + cont_two.angmom, + cont_two.angmom_components_cart, + cont_two.exps, + cont_two.coeffs, + ) + + # Return sqrt of maximum absolute value + return np.sqrt(np.max(np.abs(integral))) + + +def compute_schwarz_bounds(contractions, boys_func, compute_integral_func): + """Precompute Schwarz bounds for all shell pairs. + + Parameters + ---------- + contractions : list of GeneralizedContractionShell + List of all contracted shells. + boys_func : callable + Boys function for integral evaluation. + compute_integral_func : callable + Function to compute (ab|cd) integrals. + + Returns + ------- + bounds : np.ndarray(n_shells, n_shells) + Schwarz bounds sqrt((ab|ab)) for each shell pair. + """ + n_shells = len(contractions) + bounds = np.zeros((n_shells, n_shells)) + + for i, cont_i in enumerate(contractions): + for j in range(i, n_shells): + cont_j = contractions[j] + bounds[i, j] = compute_schwarz_bound_shell_pair( + boys_func, cont_i, cont_j, compute_integral_func + ) + bounds[j, i] = bounds[i, j] # Symmetry: (ab|ab) = (ba|ba) + + return bounds + + +def shell_pair_significant(cont_one, cont_two, threshold=1e-12): + """Check if a shell pair is significant using primitive screening. + + Uses the Gaussian product theorem: exp(-a*b/(a+b) * |A-B|^2) factor. + If this factor is below threshold for all primitive pairs, skip. + + Parameters + ---------- + cont_one : GeneralizedContractionShell + First contracted shell. + cont_two : GeneralizedContractionShell + Second contracted shell. + threshold : float + Screening threshold. + + Returns + ------- + significant : bool + True if shell pair might contribute significantly. + """ + # Distance between shell centers + r_ab_sq = np.sum((cont_one.coord - cont_two.coord) ** 2) + + if r_ab_sq < 1e-10: + # Same center, always significant + return True + + # Check if any primitive pair survives screening + for exp_a in cont_one.exps: + for exp_b in cont_two.exps: + # Gaussian decay factor + decay = np.exp(-exp_a * exp_b / (exp_a + exp_b) * r_ab_sq) + if decay > threshold: + return True + + return False + + +class SchwarzScreener: + """Class for Schwarz integral screening. + + Precomputes Schwarz bounds and provides efficient screening. + + Attributes + ---------- + bounds : np.ndarray + Schwarz bounds for all shell pairs. + threshold : float + Screening threshold. + n_screened : int + Counter for number of screened shell quartets. + n_computed : int + Counter for number of computed shell quartets. + """ + + def __init__(self, contractions, boys_func, compute_integral_func, threshold=1e-12): + """Initialize Schwarz screener. + + Parameters + ---------- + contractions : list of GeneralizedContractionShell + List of all contracted shells. + boys_func : callable + Boys function for integral evaluation. + compute_integral_func : callable + Function to compute (ab|cd) integrals. + threshold : float + Screening threshold (default: 1e-12). + """ + self.threshold = threshold + self.n_screened = 0 + self.n_computed = 0 + + # Precompute Schwarz bounds + self.bounds = compute_schwarz_bounds(contractions, boys_func, compute_integral_func) + + def is_significant(self, i, j, k, l_shell): + """Check if shell quartet (ij|kl) is significant. + + Uses Schwarz inequality: |(ij|kl)| <= sqrt((ij|ij)) * sqrt((kl|kl)) + + Parameters + ---------- + i, j, k, l_shell : int + Shell indices. + + Returns + ------- + significant : bool + True if integral might be significant, False if can be skipped. + """ + bound = self.bounds[i, j] * self.bounds[k, l_shell] + + if bound < self.threshold: + self.n_screened += 1 + return False + else: + self.n_computed += 1 + return True + + def get_statistics(self): + """Get screening statistics. + + Returns + ------- + stats : dict + Dictionary with screening statistics. + """ + total = self.n_screened + self.n_computed + if total == 0: + percent_screened = 0.0 + else: + percent_screened = 100.0 * self.n_screened / total + + return { + "n_screened": self.n_screened, + "n_computed": self.n_computed, + "total": total, + "percent_screened": percent_screened, + "speedup_factor": total / max(self.n_computed, 1), + } + + def reset_counters(self): + """Reset screening counters.""" + self.n_screened = 0 + self.n_computed = 0 diff --git a/gbasis/integrals/_two_elec_int_improved.py b/gbasis/integrals/_two_elec_int_improved.py new file mode 100644 index 00000000..6ea46fa4 --- /dev/null +++ b/gbasis/integrals/_two_elec_int_improved.py @@ -0,0 +1,732 @@ +"""Improved two-electron integrals using Obara-Saika + Head-Gordon-Pople recursions. + +This module implements the complete OS+HGP algorithm for two-electron integrals. + +Algorithm overview (full pipeline): +1. Start with Boys function F_m(T) for m = 0 to angmom_total (PR 1) +2. VRR: Build [a0|00]^m from [00|00]^m (Eq. 65) (PR 2) +3. ETR: Build [a0|c0]^0 from [a0|00]^m (Eq. 66) (PR 3) +4. Contract primitives (PR 3) +5. HRR: Build [ab|cd] from [a0|c0] (Eq. 67) (PR 4) + +References: +- Obara, S. & Saika, A. J. Chem. Phys. 1986, 84, 3963. +- Head-Gordon, M. & Pople, J. A. J. Chem. Phys. 1988, 89, 5777. +- Ahlrichs, R. Phys. Chem. Chem. Phys. 2006, 8, 3072. +""" + +import functools + +import numpy as np + +from gbasis.utils import factorial2 + + +@functools.cache +def _get_factorial2_norm(angmom_key): + """Get cached factorial2 normalization for angular momentum components. + + Parameters + ---------- + angmom_key : tuple of tuples + Angular momentum components as a tuple of tuples, e.g. + ((lx1, ly1, lz1), (lx2, ly2, lz2), ...). + + Returns + ------- + norm : np.ndarray(n,) + Normalization factors 1/sqrt(prod((2*l-1)!!)). + """ + angmom_components = np.array(angmom_key) + return 1.0 / np.sqrt(np.prod(factorial2(2 * angmom_components - 1), axis=1)) + + +def _optimized_contraction(integrals_etransf, exps, coeffs, angmoms): + """Optimized primitive contraction using einsum. + + Parameters + ---------- + integrals_etransf : np.ndarray + ETR output with shape (c_x, c_y, c_z, a_x, a_y, a_z, K_d, K_b, K_c, K_a). + exps : array-like of shape (4, K) + Primitive exponents stacked for all 4 centers (a, b, c, d). + coeffs : array-like of shape (4, K, M) + Contraction coefficients stacked for all 4 centers (a, b, c, d). + angmoms : array-like of shape (4,) + Angular momenta for all 4 centers (a, b, c, d). + + Returns + ------- + contracted : np.ndarray + Contracted integrals with shape (c_x, c_y, c_z, a_x, a_y, a_z, M_a, M_c, M_b, M_d). + """ + # Compute norms per center (supports different K per center) + norms = [((2 / np.pi) * e) ** 0.75 * (4 * e) ** (ang / 2) for e, ang in zip(exps, angmoms)] + coeffs_norm = [c * n[:, np.newaxis] for c, n in zip(coeffs, norms)] + coeffs_a_norm, coeffs_b_norm, coeffs_c_norm, coeffs_d_norm = coeffs_norm + + # Use einsum with optimization for contraction + # Input: (c_x, c_y, c_z, a_x, a_y, a_z, K_d, K_b, K_c, K_a) + # Contract one primitive at a time for memory efficiency + # K_a contraction + contracted = np.einsum("...a,aA->...A", integrals_etransf, coeffs_a_norm, optimize=True) + # K_c contraction (now axis -2 is K_c) + contracted = np.einsum("...cA,cC->...CA", contracted, coeffs_c_norm, optimize=True) + # K_b contraction (now axis -3 is K_b) + contracted = np.einsum("...bCA,bB->...CBA", contracted, coeffs_b_norm, optimize=True) + # K_d contraction (now axis -4 is K_d) + contracted = np.einsum("...dCBA,dD->...CBAD", contracted, coeffs_d_norm, optimize=True) + + # Reorder to (c_x, c_y, c_z, a_x, a_y, a_z, M_a, M_c, M_b, M_d) + # Current: (..., M_c, M_b, M_a, M_d) -> need (..., M_a, M_c, M_b, M_d) + contracted = np.moveaxis(contracted, -2, -4) + + return contracted + + +def _vertical_recursion_relation( + integrals_m, + m_max, + rel_coord_a, + coord_wac, + harm_mean, + exps_sum_one, +): + """Apply Vertical Recursion Relation (VRR) to build angular momentum on center A. + + This implements Eq. 65 from the algorithm notes: + [a+1,0|00]^m = (P-A)_i [a0|00]^m - (rho/zeta)(Q-P)_i [a0|00]^{m+1} + + a_i/(2*zeta) * ([a-1,0|00]^m - (rho/zeta)[a-1,0|00]^{m+1}) + + Parameters + ---------- + integrals_m : np.ndarray + Array containing [00|00]^m for all m values. + Shape: (m_max, K_d, K_b, K_c, K_a) + m_max : int + Maximum angular momentum order. + rel_coord_a : np.ndarray(3, K_d, K_b, K_c, K_a) + P - A coordinates for each primitive combination. + coord_wac : np.ndarray(3, K_d, K_b, K_c, K_a) + Q - P coordinates (weighted average centers difference). + harm_mean : np.ndarray(K_d, K_b, K_c, K_a) + Harmonic mean: rho = zeta*eta/(zeta+eta). + exps_sum_one : np.ndarray(K_d, K_b, K_c, K_a) + Sum of exponents: zeta = alpha_a + alpha_b. + + Returns + ------- + integrals_vert : np.ndarray + Integrals with angular momentum built on center A. + Shape: (m_max, m_max, m_max, m_max, K_d, K_b, K_c, K_a) + Axes 1, 2, 3 correspond to a_x, a_y, a_z. + """ + # Precompute coefficients for efficiency (avoid repeated division) + rho_over_zeta = harm_mean / exps_sum_one + half_over_zeta = 0.5 / exps_sum_one + + # Precompute products (avoids repeated multiplication in loops) + roz_wac_x = rho_over_zeta * coord_wac[0] + roz_wac_y = rho_over_zeta * coord_wac[1] + roz_wac_z = rho_over_zeta * coord_wac[2] + + # Initialize output array with contiguous memory + # axis 0: m, axis 1: a_x, axis 2: a_y, axis 3: a_z + integrals_vert = np.zeros((m_max, m_max, m_max, m_max, *integrals_m.shape[1:]), order="C") + + # Base case: [00|00]^m + integrals_vert[:, 0, 0, 0, ...] = integrals_m + + # VRR for x-component (a_x) + if m_max > 1: + # First step: a_x = 0 -> a_x = 1 + integrals_vert[:-1, 1, 0, 0, ...] = ( + rel_coord_a[0] * integrals_vert[:-1, 0, 0, 0, ...] + - roz_wac_x * integrals_vert[1:, 0, 0, 0, ...] + ) + # Higher a_x values (precompute a * half_over_zeta) + for a in range(1, m_max - 1): + coeff_a = a * half_over_zeta + integrals_vert[:-1, a + 1, 0, 0, ...] = ( + rel_coord_a[0] * integrals_vert[:-1, a, 0, 0, ...] + - roz_wac_x * integrals_vert[1:, a, 0, 0, ...] + + coeff_a + * ( + integrals_vert[:-1, a - 1, 0, 0, ...] + - rho_over_zeta * integrals_vert[1:, a - 1, 0, 0, ...] + ) + ) + + # VRR for y-component (a_y) + if m_max > 1: + integrals_vert[:-1, :, 1, 0, ...] = ( + rel_coord_a[1] * integrals_vert[:-1, :, 0, 0, ...] + - roz_wac_y * integrals_vert[1:, :, 0, 0, ...] + ) + for a in range(1, m_max - 1): + coeff_a = a * half_over_zeta + integrals_vert[:-1, :, a + 1, 0, ...] = ( + rel_coord_a[1] * integrals_vert[:-1, :, a, 0, ...] + - roz_wac_y * integrals_vert[1:, :, a, 0, ...] + + coeff_a + * ( + integrals_vert[:-1, :, a - 1, 0, ...] + - rho_over_zeta * integrals_vert[1:, :, a - 1, 0, ...] + ) + ) + + # VRR for z-component (a_z) + if m_max > 1: + integrals_vert[:-1, :, :, 1, ...] = ( + rel_coord_a[2] * integrals_vert[:-1, :, :, 0, ...] + - roz_wac_z * integrals_vert[1:, :, :, 0, ...] + ) + for a in range(1, m_max - 1): + coeff_a = a * half_over_zeta + integrals_vert[:-1, :, :, a + 1, ...] = ( + rel_coord_a[2] * integrals_vert[:-1, :, :, a, ...] + - roz_wac_z * integrals_vert[1:, :, :, a, ...] + + coeff_a + * ( + integrals_vert[:-1, :, :, a - 1, ...] + - rho_over_zeta * integrals_vert[1:, :, :, a - 1, ...] + ) + ) + + return integrals_vert + + +def _electron_transfer_recursion( + integrals_vert, + m_max, + m_max_c, + rel_coord_c, + rel_coord_a, + exps_sum_one, + exps_sum_two, +): + """Apply Electron Transfer Recursion (ETR) to transfer angular momentum to center C. + + This implements Eq. 66 from the algorithm notes: + [a0|c+1,0]^0 = (Q-C)_i [a0|c0]^0 + (zeta/eta)(P-A)_i [a0|c0]^0 + - (zeta/eta) [a+1,0|c0]^0 + + c_i/(2*eta) [a0|c-1,0]^0 + + a_i/(2*eta) [a-1,0|c0]^0 + + Parameters + ---------- + integrals_vert : np.ndarray + Output from VRR with angular momentum on A. + Shape: (m_max, a_x_max, a_y_max, a_z_max, K_d, K_b, K_c, K_a) + m_max : int + Maximum m value (angmom_a + angmom_b + angmom_c + angmom_d + 1). + m_max_c : int + Maximum c angular momentum (angmom_c + angmom_d + 1). + rel_coord_c : np.ndarray(3, K_d, K_b, K_c, K_a) + Q - C coordinates. + rel_coord_a : np.ndarray(3, K_d, K_b, K_c, K_a) + P - A coordinates. + exps_sum_one : np.ndarray + zeta = alpha_a + alpha_b. + exps_sum_two : np.ndarray + eta = alpha_c + alpha_d. + + Returns + ------- + integrals_etransf : np.ndarray + Integrals with angular momentum on both A and C. + Shape: (c_x_max, c_y_max, c_z_max, a_x_max, a_y_max, a_z_max, K_d, K_b, K_c, K_a) + """ + n_primitives = integrals_vert.shape[4:] + zeta_over_eta = exps_sum_one / exps_sum_two + + # Precompute coefficients (avoid repeated division in loops) + half_over_eta = 0.5 / exps_sum_two + + # Precompute combined coordinate terms for each axis + qc_plus_zoe_pa_x = rel_coord_c[0] + zeta_over_eta * rel_coord_a[0] + qc_plus_zoe_pa_y = rel_coord_c[1] + zeta_over_eta * rel_coord_a[1] + qc_plus_zoe_pa_z = rel_coord_c[2] + zeta_over_eta * rel_coord_a[2] + + # Initialize ETR output with contiguous memory + integrals_etransf = np.zeros( + (m_max_c, m_max_c, m_max_c, m_max, m_max, m_max, *n_primitives), order="C" + ) + + # Base case: discard m index (take m=0) + integrals_etransf[0, 0, 0, ...] = integrals_vert[0, ...] + + # Precompute a_indices coefficient array once + if m_max > 2: + a_coeff_x = ( + np.arange(1, m_max - 1).reshape(-1, 1, 1, *([1] * len(n_primitives))) * half_over_eta + ) + a_coeff_y = ( + np.arange(1, m_max - 1).reshape(1, 1, 1, 1, -1, 1, *([1] * len(n_primitives))) + * half_over_eta + ) + a_coeff_z = ( + np.arange(1, m_max - 1).reshape(1, 1, 1, 1, 1, -1, *([1] * len(n_primitives))) + * half_over_eta + ) + + # ETR for c_x + for c in range(m_max_c - 1): + c_coeff = c * half_over_eta # Precompute c/(2*eta) + if c == 0: + # First step: c_x = 0 -> c_x = 1 + # For a_x = 0 + integrals_etransf[1, 0, 0, 0, ...] = ( + qc_plus_zoe_pa_x * integrals_etransf[0, 0, 0, 0, ...] + - zeta_over_eta * integrals_etransf[0, 0, 0, 1, ...] + ) + # For a_x >= 1 + if m_max > 2: + integrals_etransf[1, 0, 0, 1:-1, ...] = ( + qc_plus_zoe_pa_x * integrals_etransf[0, 0, 0, 1:-1, ...] + + a_coeff_x * integrals_etransf[0, 0, 0, :-2, ...] + - zeta_over_eta * integrals_etransf[0, 0, 0, 2:, ...] + ) + else: + # General case: c_x -> c_x + 1 + integrals_etransf[c + 1, 0, 0, 0, ...] = ( + qc_plus_zoe_pa_x * integrals_etransf[c, 0, 0, 0, ...] + + c_coeff * integrals_etransf[c - 1, 0, 0, 0, ...] + - zeta_over_eta * integrals_etransf[c, 0, 0, 1, ...] + ) + if m_max > 2: + integrals_etransf[c + 1, 0, 0, 1:-1, ...] = ( + qc_plus_zoe_pa_x * integrals_etransf[c, 0, 0, 1:-1, ...] + + a_coeff_x * integrals_etransf[c, 0, 0, :-2, ...] + + c_coeff * integrals_etransf[c - 1, 0, 0, 1:-1, ...] + - zeta_over_eta * integrals_etransf[c, 0, 0, 2:, ...] + ) + + # ETR for c_y (similar structure, using precomputed coefficients) + for c in range(m_max_c - 1): + c_coeff = c * half_over_eta + if c == 0: + integrals_etransf[:, 1, 0, :, 0, ...] = ( + qc_plus_zoe_pa_y * integrals_etransf[:, 0, 0, :, 0, ...] + - zeta_over_eta * integrals_etransf[:, 0, 0, :, 1, ...] + ) + if m_max > 2: + integrals_etransf[:, 1, 0, :, 1:-1, ...] = ( + qc_plus_zoe_pa_y * integrals_etransf[:, 0, 0, :, 1:-1, ...] + + a_coeff_y * integrals_etransf[:, 0, 0, :, :-2, ...] + - zeta_over_eta * integrals_etransf[:, 0, 0, :, 2:, ...] + ) + else: + integrals_etransf[:, c + 1, 0, :, 0, ...] = ( + qc_plus_zoe_pa_y * integrals_etransf[:, c, 0, :, 0, ...] + + c_coeff * integrals_etransf[:, c - 1, 0, :, 0, ...] + - zeta_over_eta * integrals_etransf[:, c, 0, :, 1, ...] + ) + if m_max > 2: + integrals_etransf[:, c + 1, 0, :, 1:-1, ...] = ( + qc_plus_zoe_pa_y * integrals_etransf[:, c, 0, :, 1:-1, ...] + + a_coeff_y * integrals_etransf[:, c, 0, :, :-2, ...] + + c_coeff * integrals_etransf[:, c - 1, 0, :, 1:-1, ...] + - zeta_over_eta * integrals_etransf[:, c, 0, :, 2:, ...] + ) + + # ETR for c_z (similar structure, using precomputed coefficients) + for c in range(m_max_c - 1): + c_coeff = c * half_over_eta + if c == 0: + integrals_etransf[:, :, 1, :, :, 0, ...] = ( + qc_plus_zoe_pa_z * integrals_etransf[:, :, 0, :, :, 0, ...] + - zeta_over_eta * integrals_etransf[:, :, 0, :, :, 1, ...] + ) + if m_max > 2: + integrals_etransf[:, :, 1, :, :, 1:-1, ...] = ( + qc_plus_zoe_pa_z * integrals_etransf[:, :, 0, :, :, 1:-1, ...] + + a_coeff_z * integrals_etransf[:, :, 0, :, :, :-2, ...] + - zeta_over_eta * integrals_etransf[:, :, 0, :, :, 2:, ...] + ) + else: + integrals_etransf[:, :, c + 1, :, :, 0, ...] = ( + qc_plus_zoe_pa_z * integrals_etransf[:, :, c, :, :, 0, ...] + + c_coeff * integrals_etransf[:, :, c - 1, :, :, 0, ...] + - zeta_over_eta * integrals_etransf[:, :, c, :, :, 1, ...] + ) + if m_max > 2: + integrals_etransf[:, :, c + 1, :, :, 1:-1, ...] = ( + qc_plus_zoe_pa_z * integrals_etransf[:, :, c, :, :, 1:-1, ...] + + a_coeff_z * integrals_etransf[:, :, c, :, :, :-2, ...] + + c_coeff * integrals_etransf[:, :, c - 1, :, :, 1:-1, ...] + - zeta_over_eta * integrals_etransf[:, :, c, :, :, 2:, ...] + ) + + return integrals_etransf + + +def _horizontal_recursion_relation( + integrals_cont, + angmom_a, + angmom_b, + angmom_c, + angmom_d, + angmom_components_a, + angmom_components_b, + angmom_components_c, + angmom_components_d, + rel_dist_one, + rel_dist_two, +): + """Apply Horizontal Recursion Relation (HRR) to distribute angular momentum to B and D. + + This implements Eq. 67 from the algorithm notes (Head-Gordon-Pople scheme): + [ab|cd] = [a+1,0|cd] + (A-B)_i * [a,0|cd] + + In the HGP scheme, HRR is applied AFTER contraction, which is more efficient + because contracted integrals are smaller than primitive integrals. + + Parameters + ---------- + integrals_cont : np.ndarray + Contracted integrals from ETR+contraction step. + Shape: (c_x, c_y, c_z, a_x, a_y, a_z, M_a, M_c, M_b, M_d) + angmom_a/b/c/d : int + Angular momenta for each center. + angmom_components_a/b/c/d : np.ndarray(L, 3) + Angular momentum components for each center. + rel_dist_one : np.ndarray(3,) + A - B distance vector. + rel_dist_two : np.ndarray(3,) + C - D distance vector. + + Returns + ------- + integrals : np.ndarray(L_a, L_b, L_c, L_d, M_a, M_b, M_c, M_d) + Final integrals in Chemists' notation. + """ + m_max_a = angmom_a + angmom_b + 1 + m_max_c = angmom_c + angmom_d + 1 + n_cont = integrals_cont.shape[6:] # (M_a, M_c, M_b, M_d) + + # ---- HRR for center D (x and y components) ---- + integrals_horiz_d = np.zeros( + (angmom_d + 1, angmom_d + 1, m_max_c, m_max_c, m_max_c, m_max_a, m_max_a, m_max_a, *n_cont) + ) + integrals_horiz_d[0, 0, :, :, :, :, :, :, ...] = integrals_cont[ + :, :, :, :m_max_a, :m_max_a, :m_max_a, ... + ] + + # HRR x-component of d + for d in range(angmom_d): + integrals_horiz_d[d + 1, 0, :-1, :, :, :, :, :, ...] = ( + integrals_horiz_d[d, 0, 1:, :, :, :, :, :, ...] + + rel_dist_two[0] * integrals_horiz_d[d, 0, :-1, :, :, :, :, :, ...] + ) + # HRR y-component of d + for d in range(angmom_d): + integrals_horiz_d[:, d + 1, :, :-1, :, :, :, :, ...] = ( + integrals_horiz_d[:, d, :, 1:, :, :, :, :, ...] + + rel_dist_two[1] * integrals_horiz_d[:, d, :, :-1, :, :, :, :, ...] + ) + + # Select angular momentum components (x, y) for c and d, then recurse z + angmoms_c_x, angmoms_c_y, angmoms_c_z = angmom_components_c.T + angmoms_d_x, angmoms_d_y, angmoms_d_z = angmom_components_d.T + + integrals_horiz_d2 = np.zeros( + ( + angmom_d + 1, + m_max_c, + angmom_components_d.shape[0], + angmom_components_c.shape[0], + m_max_a, + m_max_a, + m_max_a, + *n_cont, + ) + ) + integrals_horiz_d2[0] = np.transpose( + integrals_horiz_d[ + angmoms_d_x.reshape(-1, 1), + angmoms_d_y.reshape(-1, 1), + angmoms_c_x.reshape(1, -1), + angmoms_c_y.reshape(1, -1), + :, + :, + :, + :, + :, + :, + :, + :, + ], + (2, 0, 1, 3, 4, 5, 6, 7, 8, 9), + ) + + # HRR z-component of d + for d in range(angmom_d): + integrals_horiz_d2[d + 1, :-1, :, :, :, :, :, ...] = ( + integrals_horiz_d2[d, 1:, :, :, :, :, :, ...] + + rel_dist_two[2] * integrals_horiz_d2[d, :-1, :, :, :, :, :, ...] + ) + + # Select z components for c and d + integrals_horiz_d2 = integrals_horiz_d2[ + angmoms_d_z.reshape(-1, 1), + angmoms_c_z.reshape(1, -1), + np.arange(angmoms_d_z.size).reshape(-1, 1), + np.arange(angmoms_c_z.size).reshape(1, -1), + :, + :, + :, + :, + :, + :, + :, + ] + + # ---- HRR for center B (x and y components) ---- + angmoms_a_x, angmoms_a_y, angmoms_a_z = angmom_components_a.T + angmoms_b_x, angmoms_b_y, angmoms_b_z = angmom_components_b.T + + integrals_horiz_b = np.zeros( + ( + angmom_b + 1, + angmom_b + 1, + m_max_a, + m_max_a, + m_max_a, + angmom_components_d.shape[0], + angmom_components_c.shape[0], + *n_cont, + ) + ) + integrals_horiz_b[0, 0, :, :, :, :, :, ...] = np.transpose( + integrals_horiz_d2[:, :, :, :, :, :, :, :, :], (2, 3, 4, 0, 1, 5, 6, 7, 8) + ) + + # HRR x-component of b + for b in range(angmom_b): + integrals_horiz_b[b + 1, 0, :-1, :, :, :, :, ...] = ( + integrals_horiz_b[b, 0, 1:, :, :, :, :, ...] + + rel_dist_one[0] * integrals_horiz_b[b, 0, :-1, :, :, :, :, ...] + ) + # HRR y-component of b + for b in range(angmom_b): + integrals_horiz_b[:, b + 1, :, :-1, :, :, :, ...] = ( + integrals_horiz_b[:, b, :, 1:, :, :, :, ...] + + rel_dist_one[1] * integrals_horiz_b[:, b, :, :-1, :, :, :, ...] + ) + + # Select angular momentum components (x, y) for a and b, then recurse z + integrals_horiz_b2 = np.zeros( + ( + angmom_b + 1, + m_max_a, + angmom_components_b.shape[0], + angmom_components_a.shape[0], + angmom_components_d.shape[0], + angmom_components_c.shape[0], + *n_cont, + ) + ) + integrals_horiz_b2[0] = np.transpose( + integrals_horiz_b[ + angmoms_b_x.reshape(-1, 1), + angmoms_b_y.reshape(-1, 1), + angmoms_a_x.reshape(1, -1), + angmoms_a_y.reshape(1, -1), + :, + :, + :, + :, + :, + :, + :, + ], + (2, 0, 1, 3, 4, 5, 6, 7, 8), + ) + + # HRR z-component of b + for b in range(angmom_b): + integrals_horiz_b2[b + 1, :-1, :, :, :, :, ...] = ( + integrals_horiz_b2[b, 1:, :, :, :, :, ...] + + rel_dist_one[2] * integrals_horiz_b2[b, :-1, :, :, :, :, ...] + ) + + # Select z components for a and b + integrals_horiz_b2 = integrals_horiz_b2[ + angmoms_b_z.reshape(-1, 1), + angmoms_a_z.reshape(1, -1), + np.arange(angmoms_b_z.size).reshape(-1, 1), + np.arange(angmoms_a_z.size).reshape(1, -1), + :, + :, + :, + :, + :, + :, + ] + + # Rearrange to final order: (L_a, L_b, L_c, L_d, M_a, M_b, M_c, M_d) + # Current: (L_b, L_a, L_d, L_c, M_a, M_c, M_b, M_d) + integrals = np.transpose(integrals_horiz_b2, (1, 0, 3, 2, 4, 6, 5, 7)) + + # Apply factorial2 normalization for angular momentum components + norm_a = _get_factorial2_norm(tuple(map(tuple, angmom_components_a))).reshape( + -1, 1, 1, 1, 1, 1, 1, 1 + ) + norm_b = _get_factorial2_norm(tuple(map(tuple, angmom_components_b))).reshape( + 1, -1, 1, 1, 1, 1, 1, 1 + ) + norm_c = _get_factorial2_norm(tuple(map(tuple, angmom_components_c))).reshape( + 1, 1, -1, 1, 1, 1, 1, 1 + ) + norm_d = _get_factorial2_norm(tuple(map(tuple, angmom_components_d))).reshape( + 1, 1, 1, -1, 1, 1, 1, 1 + ) + + integrals = integrals * norm_a * norm_b * norm_c * norm_d + + return integrals + + +def compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + angmom_a, + angmom_components_a, + exps_a, + coeffs_a, + coord_b, + angmom_b, + angmom_components_b, + exps_b, + coeffs_b, + coord_c, + angmom_c, + angmom_components_c, + exps_c, + coeffs_c, + coord_d, + angmom_d, + angmom_components_d, + exps_d, + coeffs_d, + primitive_threshold=0.0, +): + r"""Compute two-electron integrals using OS+HGP algorithm. + + This is the main entry point that combines all steps: + 1. Boys function initialization (base case [00|00]^m) + 2. VRR: Build angular momentum on center A (Eq. 65) + 3. ETR: Transfer angular momentum to center C (Eq. 66) + 4. Contract primitives (einsum-based) + 5. HRR: Distribute to centers B and D (Eq. 67, done LAST per HGP) + + Parameters + ---------- + boys_func : callable + Boys function with signature boys_func(orders, weighted_dist). + coord_a/b/c/d : np.ndarray(3,) + Centers of each contraction. + angmom_a/b/c/d : int + Angular momentum of each contraction. + angmom_components_a/b/c/d : np.ndarray(L, 3) + Angular momentum components for each contraction. + exps_a/b/c/d : np.ndarray(K,) + Primitive exponents. + coeffs_a/b/c/d : np.ndarray(K, M) + Contraction coefficients. + primitive_threshold : float, optional + Screening threshold for primitive quartets (default: 0.0, no screening). + Primitive quartets with |prefactor| < threshold are zeroed out per Eq. 64. + + Returns + ------- + integrals : np.ndarray(L_a, L_b, L_c, L_d, M_a, M_b, M_c, M_d) + Two-electron integrals in Chemists' notation. + """ + m_max = angmom_a + angmom_b + angmom_c + angmom_d + 1 + m_max_c = angmom_c + angmom_d + 1 + + # --- Pre-compute primitive quantities --- + # Reshape exponents for broadcasting: (K_d, K_b, K_c, K_a) + coord_a_4d = coord_a[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis] + coord_b_4d = coord_b[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis] + coord_c_4d = coord_c[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis] + coord_d_4d = coord_d[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis] + + exps_a_4d = exps_a[np.newaxis, np.newaxis, np.newaxis, :] + exps_b_4d = exps_b[np.newaxis, :, np.newaxis, np.newaxis] + exps_c_4d = exps_c[np.newaxis, np.newaxis, :, np.newaxis] + exps_d_4d = exps_d[:, np.newaxis, np.newaxis, np.newaxis] + + # Sum of exponents + exps_sum_one = exps_a_4d + exps_b_4d # zeta = alpha_a + alpha_b + exps_sum_two = exps_c_4d + exps_d_4d # eta = alpha_c + alpha_d + exps_sum = exps_sum_one + exps_sum_two + + # Harmonic means + harm_mean_one = (exps_a_4d * exps_b_4d) / exps_sum_one + harm_mean_two = (exps_c_4d * exps_d_4d) / exps_sum_two + harm_mean = exps_sum_one * exps_sum_two / exps_sum # rho + + # Weighted average centers + coord_wac_one = (exps_a_4d * coord_a_4d + exps_b_4d * coord_b_4d) / exps_sum_one # P + coord_wac_two = (exps_c_4d * coord_c_4d + exps_d_4d * coord_d_4d) / exps_sum_two # Q + coord_wac = coord_wac_one - coord_wac_two # P - Q + + # Relative coordinates + rel_dist_one = coord_a - coord_b # A - B (for HRR) + rel_dist_two = coord_c - coord_d # C - D (for HRR) + rel_coord_a = coord_wac_one - coord_a_4d # P - A (for VRR) + rel_coord_c = coord_wac_two - coord_c_4d # Q - C (for ETR) + + # --- Step 1: Boys function initialization --- + weighted_dist = harm_mean * np.sum(coord_wac**2, axis=0) + prefactor = ( + (2 * np.pi**2.5) + / (exps_sum_one * exps_sum_two * exps_sum**0.5) + * np.exp(-harm_mean_one * np.sum((coord_a_4d - coord_b_4d) ** 2, axis=0)) + * np.exp(-harm_mean_two * np.sum((coord_c_4d - coord_d_4d) ** 2, axis=0)) + ) + + # --- Primitive-level screening (Eq. 64) --- + if primitive_threshold > 0: + prefactor = np.where(np.abs(prefactor) >= primitive_threshold, prefactor, 0.0) + + orders = np.arange(m_max)[:, None, None, None, None] + integrals_m = boys_func(orders, weighted_dist[None, :, :, :, :], rho=harm_mean) * prefactor + + # --- Step 2: VRR --- + integrals_vert = _vertical_recursion_relation( + integrals_m, m_max, rel_coord_a, coord_wac, harm_mean, exps_sum_one + ) + + # --- Step 3: ETR --- + integrals_etransf = _electron_transfer_recursion( + integrals_vert, m_max, m_max_c, rel_coord_c, rel_coord_a, exps_sum_one, exps_sum_two + ) + + # --- Step 4: Contract primitives --- + integrals_cont = _optimized_contraction( + integrals_etransf, + (exps_a, exps_b, exps_c, exps_d), + (coeffs_a, coeffs_b, coeffs_c, coeffs_d), + (angmom_a, angmom_b, angmom_c, angmom_d), + ) + + # --- Step 5: HRR (done LAST per HGP scheme) --- + integrals = _horizontal_recursion_relation( + integrals_cont, + angmom_a, + angmom_b, + angmom_c, + angmom_d, + angmom_components_a, + angmom_components_b, + angmom_components_c, + angmom_components_d, + rel_dist_one, + rel_dist_two, + ) + + return integrals diff --git a/gbasis/integrals/boys_functions.py b/gbasis/integrals/boys_functions.py new file mode 100644 index 00000000..f9cb8ac6 --- /dev/null +++ b/gbasis/integrals/boys_functions.py @@ -0,0 +1,350 @@ +"""Boys functions for two-electron interaction potentials. + +This module implements the Boys function, which is the starting point for +computing two-electron integrals using the Obara-Saika recursion relations. + +The standard Boys function is defined as: + + F_m(T) = integral from 0 to 1 of t^(2m) * exp(-T*t^2) dt + +For non-Coulomb potentials (erf-attenuated, erfc-attenuated), the modified +Boys functions follow from Ahlrichs (2006), where the key modification is +a rescaling of the argument T and a prefactor depending on omega and rho. + +References: +- Helgaker, T.; Jorgensen, P.; Olsen, J. "Molecular Electronic-Structure Theory" + (eq. 9.8.39 for hyp1f1 representation) +- Ahlrichs, R. "A simple algebraic derivation of the Obara-Saika scheme for + general two-electron interaction potentials." + Phys. Chem. Chem. Phys., 2006, 8, 3072-3077. +""" + +import numpy as np +from scipy.special import hyp1f1 + + +def boys_function_standard(orders, weighted_dist): + r"""Compute standard Boys function for Coulomb potential (1/r12). + + The Coulombic Boys function is defined as: + + .. math:: + + F_m(T) = \int_0^1 t^{2m} e^{-T t^2} dt + + This can be expressed in terms of the Kummer confluent hypergeometric + function (hyp1f1) as shown in Helgaker (eq. 9.8.39). + + Parameters + ---------- + orders : np.ndarray + Differentiation order of the Boys function (m values). + Shape can be (M,) or (M, 1, 1, ...) for broadcasting. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + T = rho * |P - Q|^2 where rho = zeta*eta/(zeta+eta). + + Returns + ------- + boys_eval : np.ndarray + Boys function values F_m(T) with shape determined by broadcasting + of orders and weighted_dist. + + Notes + ----- + For the Coulomb potential g(r) = 1/r: + - G_0(rho, T) = (2*pi/rho) * F_0(T) + + This is the standard case used for electron repulsion integrals. + + There's some documented instability for hyp1f1, mainly for large values. + For typical quantum chemistry calculations (m < 20), the values are stable. + """ + return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) + + +def boys_function_all_orders(m_max, weighted_dist): + r"""Compute Boys function for all orders from 0 to m_max. + + Returns F_m(T) for m = 0, 1, ..., m_max simultaneously using + scipy's hyp1f1, which is numerically stable for all practical T + and m values encountered in quantum chemistry. + + Parameters + ---------- + m_max : int + Maximum order of the Boys function needed. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + + Returns + ------- + boys_all : np.ndarray + Boys function values for all orders from 0 to m_max. + Shape: (m_max + 1, *weighted_dist.shape) + """ + # Vectorize hyp1f1 across all orders with broadcasting to avoid + # per-order Python loops while keeping numerical behaviour identical + # to the standard expression. hyp1f1 is stable for the practical + # range m < 20 and T encountered here. + + T = np.asarray(weighted_dist) + orders = np.arange(m_max + 1, dtype=np.result_type(T, np.float64)) + + # Reshape orders to broadcast over T's dimensions: (m_max+1, 1, 1, ...) + orders_shape = (m_max + 1,) + (1,) * T.ndim + orders_b = orders.reshape(orders_shape) + + boys_vals = hyp1f1(orders_b + 0.5, orders_b + 1.5, -T) / (2 * orders_b + 1) + return boys_vals + + +def boys_function_recursion(m_max, weighted_dist): + r"""Compute Boys function for all orders using downward recursion. + + Computes F_m(T) for m = 0, 1, ..., m_max using the stable downward + recursion (Eq. 71 from algorithm notes): + + .. math:: + + F_m(T) = \frac{2T \cdot F_{m+1}(T) + e^{-T}}{2m + 1} + + The highest order F_{m_max}(T) is computed via hyp1f1 as the starting + value, then all lower orders are obtained by recursing downward. + This is more efficient than calling hyp1f1 for each order separately. + + Parameters + ---------- + m_max : int + Maximum order of the Boys function needed. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + + Returns + ------- + boys_all : np.ndarray + Boys function values for all orders from 0 to m_max. + Shape: (m_max + 1, *weighted_dist.shape) + """ + T = np.asarray(weighted_dist, dtype=np.float64) + result = np.empty((m_max + 1, *T.shape), dtype=np.float64) + + # Starting value: compute F_{m_max}(T) via hyp1f1 + result[m_max] = hyp1f1(m_max + 0.5, m_max + 1.5, -T) / (2 * m_max + 1) + + # Precompute exp(-T) once + exp_neg_T = np.exp(-T) + + # Downward recursion: F_m = (2T * F_{m+1} + exp(-T)) / (2m + 1) + for m in range(m_max - 1, -1, -1): + result[m] = (2.0 * T * result[m + 1] + exp_neg_T) / (2 * m + 1) + + return result + + +def boys_function_erf(orders, weighted_dist, rho, omega): + r"""Compute Boys function for erf-attenuated Coulomb potential. + + For the potential g(r) = erf(omega * r) / r, the modified Boys function is: + + .. math:: + + G_m(\rho, T) = \frac{2\pi}{\rho} \left(\frac{\omega^2}{\omega^2 + \rho}\right)^{m+1/2} + F_m\left(\frac{\omega^2 T}{\omega^2 + \rho}\right) + + This is used in range-separated DFT for the long-range part of the + electron-electron interaction. + + Parameters + ---------- + orders : np.ndarray + Differentiation orders (m values). + weighted_dist : np.ndarray + Weighted interatomic distances (T = rho * |P-Q|^2). + rho : np.ndarray + Harmonic mean of exponent sums: rho = zeta*eta/(zeta+eta). + omega : float + Range-separation parameter. + + Returns + ------- + boys_eval : np.ndarray + Modified Boys function values. + + Notes + ----- + As omega -> infinity, erf(omega*r)/r -> 1/r, recovering the standard Coulomb. + As omega -> 0, erf(omega*r)/r -> 0, the interaction vanishes. + """ + scaling = omega**2 / (omega**2 + rho) + T_modified = scaling * weighted_dist + return scaling ** (orders + 0.5) * boys_function_standard(orders, T_modified) + + +def boys_function_erfc(orders, weighted_dist, rho, omega): + r"""Compute Boys function for erfc-attenuated Coulomb potential. + + For the potential g(r) = erfc(omega * r) / r, which is the short-range + complement of the erf-attenuated potential: + + .. math:: + + G_m^{erfc} = G_m^{Coulomb} - G_m^{erf} + + This is used in range-separated DFT for the short-range part of the + electron-electron interaction. + + Parameters + ---------- + orders : np.ndarray + Differentiation orders (m values). + weighted_dist : np.ndarray + Weighted interatomic distances (T = rho * |P-Q|^2). + rho : np.ndarray + Harmonic mean of exponent sums: rho = zeta*eta/(zeta+eta). + omega : float + Range-separation parameter. + + Returns + ------- + boys_eval : np.ndarray + Modified Boys function values. + + Notes + ----- + erfc(omega*r)/r = 1/r - erf(omega*r)/r, so: + G_m^{erfc} = F_m(T) - (omega^2/(omega^2+rho))^(m+0.5) * F_m(omega^2*T/(omega^2+rho)) + + As omega -> infinity, erfc(omega*r)/r -> 0 (short-range vanishes). + As omega -> 0, erfc(omega*r)/r -> 1/r (full Coulomb). + """ + return boys_function_standard(orders, weighted_dist) - boys_function_erf( + orders, weighted_dist, rho, omega + ) + + +def boys_function_mpmath(m_max, weighted_dist, dps=50): + r"""Compute Boys function using mpmath for high-precision reference values. + + Uses mpmath.hyp1f1 with arbitrary precision arithmetic. Useful as a + reference for validating other implementations, or for edge cases where + scipy's hyp1f1 may lose precision (very high m or specific T ranges). + + Parameters + ---------- + m_max : int + Maximum order of the Boys function needed. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + dps : int + Decimal places of precision for mpmath (default: 50). + + Returns + ------- + boys_all : np.ndarray + Boys function values for all orders from 0 to m_max. + Shape: (m_max + 1, *weighted_dist.shape) + + Raises + ------ + ImportError + If mpmath is not installed. + + Notes + ----- + This function is SLOW compared to hyp1f1 or recursion-based methods. + It is intended only for validation and reference, not for production use. + Install mpmath via: pip install mpmath + """ + try: + import mpmath + except ImportError as err: + raise ImportError( + "mpmath is required for boys_function_mpmath. Install it with: pip install mpmath" + ) from err + + T = np.asarray(weighted_dist, dtype=np.float64) + result = np.empty((m_max + 1, *T.shape), dtype=np.float64) + + old_dps = mpmath.mp.dps + mpmath.mp.dps = dps + try: + for idx in np.ndindex(T.shape): + t_val = mpmath.mpf(float(T[idx])) + for m in range(m_max + 1): + val = mpmath.hyp1f1(m + 0.5, m + 1.5, -t_val) / (2 * m + 1) + result[(m, *idx)] = float(val) + finally: + mpmath.mp.dps = old_dps + + return result + + +def get_boys_function(potential="coulomb", omega=None): + """Get the appropriate Boys function for a given potential type. + + Parameters + ---------- + potential : str + Type of two-electron potential. Options: + - "coulomb" or "standard" or "1/r": Standard 1/r12 Coulomb potential + - "erf" or "erf_coulomb": erf-attenuated (long-range) potential + - "erfc" or "erfc_coulomb": erfc-attenuated (short-range) potential + omega : float, optional + Range-separation parameter for erf/erfc potentials. + Required when potential is "erf" or "erfc". + + Returns + ------- + boys_func : callable + Boys function with signature boys_func(orders, weighted_dist, rho=None). + + Raises + ------ + ValueError + If potential type is not recognized. + If omega is not provided for erf/erfc potentials. + + Examples + -------- + >>> boys = get_boys_function("coulomb") + >>> F0 = boys(np.array([0]), np.array([1.0])) + + >>> boys_erf = get_boys_function("erf", omega=0.4) + >>> F0_erf = boys_erf(np.array([0]), np.array([1.0]), rho=0.5) + """ + potential = potential.lower() + + if potential in ["coulomb", "standard", "1/r"]: + + def boys_func(orders, weighted_dist, rho=None): + return boys_function_standard(orders, weighted_dist) + + return boys_func + + elif potential in ["erf", "erf_coulomb"]: + if omega is None: + raise ValueError("omega parameter is required for erf-attenuated potential") + + def boys_func(orders, weighted_dist, rho=None): + if rho is None: + raise ValueError("rho parameter is required for erf-attenuated potential") + return boys_function_erf(orders, weighted_dist, rho, omega) + + return boys_func + + elif potential in ["erfc", "erfc_coulomb"]: + if omega is None: + raise ValueError("omega parameter is required for erfc-attenuated potential") + + def boys_func(orders, weighted_dist, rho=None): + if rho is None: + raise ValueError("rho parameter is required for erfc-attenuated potential") + return boys_function_erfc(orders, weighted_dist, rho, omega) + + return boys_func + + else: + raise ValueError( + f"Unknown potential type: {potential}. " f"Valid options: coulomb, erf, erfc" + ) diff --git a/gbasis/integrals/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index 4643bbbe..295b0cfa 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -1,12 +1,20 @@ """Electron-electron repulsion integral.""" + +import numpy as np + from gbasis.base_four_symm import BaseFourIndexSymmetric from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals._schwarz_screening import SchwarzScreener from gbasis.integrals._two_elec_int import ( _compute_two_elec_integrals, _compute_two_elec_integrals_angmom_zero, ) +from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp +from gbasis.integrals.boys_functions import get_boys_function from gbasis.integrals.point_charge import PointChargeIntegral -import numpy as np + +# Cache for Schwarz screeners keyed by (basis_ids, threshold) +_schwarz_screener_cache = {} class ElectronRepulsionIntegral(BaseFourIndexSymmetric): @@ -271,3 +279,226 @@ def electron_repulsion_integral(basis, transform=None, notation="physicist"): if notation == "physicist": array = np.transpose(array, (0, 2, 1, 3)) return array + + +class ElectronRepulsionIntegralImproved(BaseFourIndexSymmetric): + """Class for constructing electron-electron repulsion integrals using OS+HGP algorithm. + + This class uses the improved Obara-Saika + Head-Gordon-Pople recursion scheme, + which applies the Horizontal Recursion Relation (HRR) after contraction for + better computational efficiency. + + The first four axes of the returned array are associated with the given set of contracted + Gaussian (or a linear combination of a set of Gaussians). + + Attributes + ---------- + _axes_contractions : tuple of tuple of GeneralizedContractionShell + Sets of contractions associated with each axis of the array. + contractions : tuple of GeneralizedContractionShell + Contractions that are associated with the four indices of the array. + + """ + + boys_func = get_boys_function("coulomb") + + @classmethod + def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four, **kwargs): + r"""Return electron-electron repulsion integral using the OS+HGP algorithm. + + Parameters + ---------- + cont_one : GeneralizedContractionShell + Contracted Cartesian Gaussians (of the same shell) associated with the first index. + cont_two : GeneralizedContractionShell + Contracted Cartesian Gaussians (of the same shell) associated with the second index. + cont_three : GeneralizedContractionShell + Contracted Cartesian Gaussians (of the same shell) associated with the third index. + cont_four : GeneralizedContractionShell + Contracted Cartesian Gaussians (of the same shell) associated with the fourth index. + + Returns + ------- + array_cont : np.ndarray(M_1, L_cart_1, M_2, L_cart_2, M_3, L_cart_3, M_4, L_cart_4) + Electron-electron repulsion integral associated with the given contractions. + Integrals are in Chemists' notation. + + Raises + ------ + TypeError + If any contraction is not a `GeneralizedContractionShell` instance. + + """ + if not isinstance(cont_one, GeneralizedContractionShell): + raise TypeError("`cont_one` must be a `GeneralizedContractionShell` instance.") + if not isinstance(cont_two, GeneralizedContractionShell): + raise TypeError("`cont_two` must be a `GeneralizedContractionShell` instance.") + if not isinstance(cont_three, GeneralizedContractionShell): + raise TypeError("`cont_three` must be a `GeneralizedContractionShell` instance.") + if not isinstance(cont_four, GeneralizedContractionShell): + raise TypeError("`cont_four` must be a `GeneralizedContractionShell` instance.") + + # --- Schwarz screening: skip negligible shell quartets --- + screener = kwargs.get("screener") + index_map = kwargs.get("contraction_index_map") + if screener is not None and index_map is not None: + i = index_map.get(id(cont_one)) + j = index_map.get(id(cont_two)) + k = index_map.get(id(cont_three)) + l_idx = index_map.get(id(cont_four)) + if i is not None and j is not None and k is not None and l_idx is not None: + if not screener.is_significant(i, j, k, l_idx): + shape = ( + cont_one.coeffs.shape[1], + len(cont_one.angmom_components_cart), + cont_two.coeffs.shape[1], + len(cont_two.angmom_components_cart), + cont_three.coeffs.shape[1], + len(cont_three.angmom_components_cart), + cont_four.coeffs.shape[1], + len(cont_four.angmom_components_cart), + ) + return np.zeros(shape) + + # --- Contraction reordering for efficiency (l_a >= l_b, l_c >= l_d, l_a >= l_c) --- + bra_swapped = cont_one.angmom < cont_two.angmom + if bra_swapped: + cont_one, cont_two = cont_two, cont_one + + ket_swapped = cont_three.angmom < cont_four.angmom + if ket_swapped: + cont_three, cont_four = cont_four, cont_three + + braket_swapped = cont_one.angmom < cont_three.angmom + if braket_swapped: + cont_one, cont_three = cont_three, cont_one + cont_two, cont_four = cont_four, cont_two + + integrals = compute_two_electron_integrals_os_hgp( + cls.boys_func, + cont_one.coord, + cont_one.angmom, + cont_one.angmom_components_cart, + cont_one.exps, + cont_one.coeffs, + cont_two.coord, + cont_two.angmom, + cont_two.angmom_components_cart, + cont_two.exps, + cont_two.coeffs, + cont_three.coord, + cont_three.angmom, + cont_three.angmom_components_cart, + cont_three.exps, + cont_three.coeffs, + cont_four.coord, + cont_four.angmom, + cont_four.angmom_components_cart, + cont_four.exps, + cont_four.coeffs, + ) + + # --- Un-swap axes to restore original ordering --- + # Output shape from compute: (L_a, L_b, L_c, L_d, M_a, M_b, M_c, M_d) + if braket_swapped: + integrals = np.transpose(integrals, (2, 3, 0, 1, 6, 7, 4, 5)) + if ket_swapped: + integrals = np.swapaxes(np.swapaxes(integrals, 2, 3), 6, 7) + if bra_swapped: + integrals = np.swapaxes(np.swapaxes(integrals, 0, 1), 4, 5) + + integrals = np.transpose(integrals, (4, 0, 5, 1, 6, 2, 7, 3)) + + return integrals + + +def electron_repulsion_integral_improved( + basis, transform=None, notation="physicist", schwarz_threshold=0.0 +): + r"""Return the electron repulsion integrals using the improved OS+HGP algorithm. + + This function uses the Obara-Saika + Head-Gordon-Pople recursion scheme, + which is more efficient than the original implementation because HRR is + applied after contraction. + + In the Chemists' notation, the integrals are: + + .. math:: + + \int \phi^*_a(\mathbf{r}_1) \phi_b(\mathbf{r}_1) + \frac{1}{|\mathbf{r}_1 - \mathbf{r}_2|} + \phi^*_c(\mathbf{r}_2) \phi_d(\mathbf{r}_2) d\mathbf{r} + + And in the Physicists' notation: + + .. math:: + + \int \phi^*_a(\mathbf{r}_1) \phi^*_b(\mathbf{r}_2) + \frac{1}{|\mathbf{r}_1 - \mathbf{r}_2|} + \phi_c(\mathbf{r}_1) \phi_d(\mathbf{r}_2) d\mathbf{r} + + Parameters + ---------- + basis : list/tuple of GeneralizedContractionShell + Shells of generalized contractions. + transform : np.ndarray(K, K_cont) + Transformation matrix from the basis set in the given coordinate system (e.g. AO) to linear + combinations of contractions (e.g. MO). + Default is no transformation. + notation : {"physicist", "chemist"} + Convention with which the integrals are ordered. + Default is Physicists' notation. + schwarz_threshold : float + Schwarz screening threshold. Shell quartets with Schwarz bound below + this threshold are skipped. Default is 0.0 (no screening). + + Returns + ------- + array : np.ndarray(K, K, K, K) + Electron-electron repulsion integral of the given basis set. + + Raises + ------ + ValueError + If `notation` is not one of "physicist" or "chemist". + + """ + if notation not in ["physicist", "chemist"]: + raise ValueError("`notation` must be one of 'physicist' or 'chemist'") + + screening_kwargs = {} + if schwarz_threshold > 0: + cache_key = (tuple(id(shell) for shell in basis), schwarz_threshold) + if cache_key not in _schwarz_screener_cache: + _schwarz_screener_cache[cache_key] = SchwarzScreener( + list(basis), + get_boys_function("coulomb"), + compute_two_electron_integrals_os_hgp, + schwarz_threshold, + ) + screener = _schwarz_screener_cache[cache_key] + index_map = {id(shell): i for i, shell in enumerate(basis)} + screening_kwargs = {"screener": screener, "contraction_index_map": index_map} + + coord_type = [ct for ct in [shell.coord_type for shell in basis]] + + if transform is not None: + array = ElectronRepulsionIntegralImproved(basis).construct_array_lincomb( + transform, coord_type, **screening_kwargs + ) + elif all(ct == "cartesian" for ct in coord_type): + array = ElectronRepulsionIntegralImproved(basis).construct_array_cartesian( + **screening_kwargs + ) + elif all(ct == "spherical" for ct in coord_type): + array = ElectronRepulsionIntegralImproved(basis).construct_array_spherical( + **screening_kwargs + ) + else: + array = ElectronRepulsionIntegralImproved(basis).construct_array_mix( + coord_type, **screening_kwargs + ) + + if notation == "physicist": + array = np.transpose(array, (0, 2, 1, 3)) + return array diff --git a/gbasis/integrals/point_charge.py b/gbasis/integrals/point_charge.py index 294f1bd2..fe2a7cd9 100644 --- a/gbasis/integrals/point_charge.py +++ b/gbasis/integrals/point_charge.py @@ -1,9 +1,11 @@ """Module for computing point charge integrals.""" + +import numpy as np + from gbasis.base_two_symm import BaseTwoIndexSymmetric from gbasis.contractions import GeneralizedContractionShell from gbasis.integrals._one_elec_int import _compute_one_elec_integrals -import numpy as np -from scipy.special import hyp1f1 # pylint: disable=E0611 +from gbasis.integrals.boys_functions import boys_function_standard class PointChargeIntegral(BaseTwoIndexSymmetric): @@ -65,7 +67,7 @@ class PointChargeIntegral(BaseTwoIndexSymmetric): """ @staticmethod - def boys_func(orders, weighted_dist): + def boys_func(orders, weighted_dist, rho=None): r"""Return the value of Boys function for the given orders and weighted distances. The Coulombic Boys function can be written as a renormalized special case of the Kummer @@ -113,7 +115,7 @@ def boys_func(orders, weighted_dist): side. """ - return hyp1f1(orders + 1 / 2, orders + 3 / 2, -weighted_dist) / (2 * orders + 1) + return boys_function_standard(orders, weighted_dist) @classmethod def construct_array_contraction( diff --git a/notebooks/tutorial/Tutorial 3 - Two-Electron Integral Optimizations.ipynb b/notebooks/tutorial/Tutorial 3 - Two-Electron Integral Optimizations.ipynb new file mode 100644 index 00000000..c49e25b9 --- /dev/null +++ b/notebooks/tutorial/Tutorial 3 - Two-Electron Integral Optimizations.ipynb @@ -0,0 +1,546 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# GBasis Tutorial - Two-Electron Integral Optimizations\n", + "\n", + "This tutorial covers the **optimized two-electron integral computation** in GBasis, including:\n", + "\n", + "1. **OS+HGP Algorithm** - Obara-Saika + Head-Gordon-Pople recursion scheme\n", + "2. **Boys Functions** - Standard Coulomb and range-separated (erf/erfc) potentials\n", + "3. **Schwarz Screening** - Skip negligible integrals for speedup\n", + "4. **High Angular Momentum** - Stable computation for d/f orbitals\n", + "5. **PySCF Comparison** - Validation against a trusted reference\n", + "\n", + "These optimizations are implemented as part of **Issue #221** to improve 2-electron integral performance." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Setup and Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import time\n", + "\n", + "# GBasis imports\n", + "from gbasis.contractions import GeneralizedContractionShell\n", + "from gbasis.integrals.electron_repulsion import (\n", + " ElectronRepulsionIntegral,\n", + " ElectronRepulsionIntegralImproved,\n", + " electron_repulsion_integral,\n", + " electron_repulsion_integral_improved,\n", + ")\n", + "\n", + "# Low-level OS+HGP algorithm\n", + "from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp\n", + "from gbasis.integrals.point_charge import PointChargeIntegral\n", + "\n", + "# Schwarz screening\n", + "from gbasis.integrals._screening import (\n", + " SchwarzScreener,\n", + " shell_pair_significant,\n", + " compute_schwarz_bounds,\n", + ")\n", + "\n", + "# Boys functions for different potentials\n", + "from gbasis.integrals.boys_functions import (\n", + " boys_function_standard,\n", + " boys_function_erf,\n", + " boys_function_erfc,\n", + " get_boys_function,\n", + ")\n", + "\n", + "print(\"All imports successful!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Creating Test Molecules\n", + "\n", + "We'll create hydrogen chain molecules to demonstrate the optimizations.\n", + "Larger separations between atoms lead to more sparsity and better screening." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def make_shell(coord, angmom, exps, coeffs):\n", + " \"\"\"Create a GeneralizedContractionShell.\"\"\"\n", + " return GeneralizedContractionShell(\n", + " angmom,\n", + " np.array(coord),\n", + " np.array(coeffs).reshape(-1, 1),\n", + " np.array(exps),\n", + " 'cartesian'\n", + " )\n", + "\n", + "def make_h_chain(n_atoms, spacing):\n", + " \"\"\"Create a hydrogen chain with STO-3G basis.\"\"\"\n", + " # STO-3G exponents and coefficients for H\n", + " exps = [3.42525091, 0.62391373, 0.16885540]\n", + " coeffs = [0.15432897, 0.53532814, 0.44463454]\n", + " \n", + " basis = []\n", + " for i in range(n_atoms):\n", + " coord = [i * spacing, 0.0, 0.0]\n", + " basis.append(make_shell(coord, 0, exps, coeffs))\n", + " \n", + " return basis\n", + "\n", + "# Create test molecules\n", + "h4_compact = make_h_chain(4, spacing=1.4) # Compact: H-H = 1.4 Bohr\n", + "h6_extended = make_h_chain(6, spacing=15.0) # Extended: H-H = 15 Bohr\n", + "\n", + "print(f\"Compact H4: {len(h4_compact)} shells, 1.4 Bohr spacing\")\n", + "print(f\"Extended H6: {len(h6_extended)} shells, 15.0 Bohr spacing\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. The OS+HGP Algorithm\n", + "\n", + "The improved implementation uses the **Obara-Saika + Head-Gordon-Pople** algorithm:\n", + "\n", + "```\n", + "Step 1: Boys Function F_m(T) = initial values\n", + "Step 2: VRR (Vertical) [a0|00]^m from [00|00]^m\n", + "Step 3: ETR (Electron) [a0|c0]^0 from [a0|00]^m \n", + "Step 4: Contract Sum over primitives (HGP: do early!)\n", + "Step 5: HRR (Horizontal) [ab|cd] from [a0|c0] (HGP: do last!)\n", + "```\n", + "\n", + "### Key Optimization (HGP):\n", + "The **Head-Gordon-Pople** optimization does HRR **after** contraction.\n", + "This is more efficient because contracted integrals are smaller than primitive integrals." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(\"\"\"\n", + "Obara-Saika + Head-Gordon-Pople Algorithm\n", + "==========================================\n", + "\n", + "Traditional: HGP Optimized:\n", + "------------ --------------\n", + "Boys F_m(T) Boys F_m(T)\n", + " | |\n", + " VRR VRR\n", + " | |\n", + " ETR ETR\n", + " | |\n", + " HRR <-- expensive! Contract <-- do this first!\n", + " | |\n", + " Contract HRR <-- now cheaper!\n", + " | |\n", + " [ab|cd] [ab|cd]\n", + "\n", + "Why is HGP better?\n", + "- HRR has NO 'm' index (no auxiliary function dependence)\n", + "- After contraction, fewer terms to process\n", + "- Especially beneficial for contracted basis sets (like cc-pVTZ)\n", + "\"\"\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Comparing Original vs Improved Implementation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Compare timing: original vs improved\n", + "basis = make_h_chain(4, spacing=1.4)\n", + "\n", + "# Original implementation\n", + "start = time.perf_counter()\n", + "eri_old = electron_repulsion_integral(basis, notation=\"chemist\")\n", + "time_old = time.perf_counter() - start\n", + "\n", + "# Improved OS+HGP implementation\n", + "start = time.perf_counter()\n", + "eri_new = electron_repulsion_integral_improved(basis, notation=\"chemist\")\n", + "time_new = time.perf_counter() - start\n", + "\n", + "# Verify they match\n", + "max_diff = np.max(np.abs(eri_old - eri_new))\n", + "\n", + "print(f\"Original implementation: {time_old:.4f} seconds\")\n", + "print(f\"Improved (OS+HGP): {time_new:.4f} seconds\")\n", + "print(f\"Max difference: {max_diff:.2e}\")\n", + "print(f\"Results match: {np.allclose(eri_old, eri_new, rtol=1e-10)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Boys Functions for Different Potentials\n", + "\n", + "GBasis supports various two-electron potentials beyond the standard Coulomb $1/r_{12}$:\n", + "\n", + "| Potential | Formula | Use Case |\n", + "|-----------|---------|----------|\n", + "| **Coulomb** | $1/r_{12}$ | Standard electron repulsion |\n", + "| **erf (long-range)** | $\\text{erf}(\\omega r)/r$ | Range-separated hybrids (LC-) |\n", + "| **erfc (short-range)** | $\\text{erfc}(\\omega r)/r$ | Short-range exchange |\n", + "\n", + "The key identity: $\\text{erf}(\\omega r)/r + \\text{erfc}(\\omega r)/r = 1/r$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Compare Boys functions for different potentials\n", + "T_values = np.array([0.0, 0.5, 1.0, 2.0, 5.0])\n", + "m = np.array([0]) # F_0(T)\n", + "omega = 0.5 # Range-separation parameter\n", + "\n", + "print(\"Boys Function F_0(T) for Different Potentials:\")\n", + "print(\"-\" * 50)\n", + "print(f\"{'T':>8} {'Coulomb':>12} {'erf':>12} {'erfc':>12}\")\n", + "print(\"-\" * 50)\n", + "\n", + "for t in T_values:\n", + " t_arr = np.array([t])\n", + " rho = np.array([0.8])\n", + " \n", + " coulomb = boys_function_standard(m, t_arr)[0]\n", + " erf_val = boys_function_erf(m, t_arr, rho, omega)[0]\n", + " erfc_val = boys_function_erfc(m, t_arr, rho, omega)[0]\n", + " \n", + " print(f\"{t:>8.1f} {coulomb:>12.6f} {erf_val:>12.6f} {erfc_val:>12.6f}\")\n", + "\n", + "print(\"\\nNote: erf + erfc = Coulomb (range separation identity)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify: erf + erfc = Coulomb\n", + "T_test = np.array([1.0])\n", + "rho = np.array([0.8])\n", + "\n", + "coulomb = boys_function_standard(m, T_test)[0]\n", + "erf_val = boys_function_erf(m, T_test, rho, omega)[0]\n", + "erfc_val = boys_function_erfc(m, T_test, rho, omega)[0]\n", + "\n", + "print(\"Range Separation Verification:\")\n", + "print(f\" Coulomb: {coulomb:.10f}\")\n", + "print(f\" erf + erfc: {erf_val + erfc_val:.10f}\")\n", + "print(f\" Difference: {abs(coulomb - (erf_val + erfc_val)):.2e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Using the `get_boys_function` Factory" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Factory function for different potentials\n", + "print(\"Using get_boys_function() factory:\")\n", + "print(\"-\" * 40)\n", + "\n", + "# Standard Coulomb\n", + "boys_coulomb = get_boys_function(\"coulomb\")\n", + "print(f\"Coulomb: F_0(1.0) = {boys_coulomb(np.array([0]), np.array([1.0]))[0]:.6f}\")\n", + "\n", + "# Long-range (erf)\n", + "boys_lr = get_boys_function(\"erf\", omega=0.5)\n", + "print(f\"Long-range (omega=0.5): F_0(1.0) = {boys_lr(np.array([0]), np.array([1.0]), rho=np.array([0.8]))[0]:.6f}\")\n", + "\n", + "# Short-range (erfc)\n", + "boys_sr = get_boys_function(\"erfc\", omega=0.5)\n", + "print(f\"Short-range (omega=0.5): F_0(1.0) = {boys_sr(np.array([0]), np.array([1.0]), rho=np.array([0.8]))[0]:.6f}\")\n", + "\n", + "print(\"\\nAvailable potentials: 'coulomb', 'erf', 'erfc'\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Schwarz Screening\n", + "\n", + "### What is Schwarz Screening?\n", + "\n", + "The **Schwarz inequality** states:\n", + "\n", + "$$|(ab|cd)| \\leq \\sqrt{(ab|ab)} \\times \\sqrt{(cd|cd)}$$\n", + "\n", + "If this upper bound is smaller than a threshold (e.g., $10^{-12}$), we can skip the integral computation entirely!\n", + "\n", + "### Benefits:\n", + "- For extended molecules, 60-90% of integrals can be skipped\n", + "- Speedup of 2-10x depending on system size and geometry" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Demonstrate Schwarz screening on extended system\n", + "boys_func = PointChargeIntegral.boys_func\n", + "\n", + "print(\"=\" * 55)\n", + "print(\"Schwarz Screening Comparison\")\n", + "print(\"=\" * 55)\n", + "\n", + "for name, basis, spacing in [\n", + " (\"Compact H4 (1.4 Bohr)\", h4_compact, 1.4),\n", + " (\"Extended H6 (15.0 Bohr)\", h6_extended, 15.0),\n", + "]:\n", + " screener = SchwarzScreener(\n", + " basis, boys_func, compute_two_electron_integrals_os_hgp\n", + " )\n", + " \n", + " n = len(basis)\n", + " for i in range(n):\n", + " for j in range(n):\n", + " for k in range(n):\n", + " for l in range(n):\n", + " screener.is_significant(i, j, k, l)\n", + " \n", + " stats = screener.get_statistics()\n", + " \n", + " print(f\"\\n{name}:\")\n", + " print(f\" Total shell quartets: {stats['total']}\")\n", + " print(f\" Computed: {stats['n_computed']}\")\n", + " print(f\" Screened (skipped): {stats['n_screened']}\")\n", + " print(f\" Percent screened: {stats['percent_screened']:.1f}%\")\n", + " print(f\" Speedup factor: {stats['speedup_factor']:.2f}x\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Shell-Pair Prescreening\n", + "\n", + "Before even computing Schwarz bounds, we can use the **Gaussian product theorem** to\n", + "quickly check if a shell pair is significant:\n", + "\n", + "$$\\exp\\left(-\\frac{\\alpha_a \\alpha_b}{\\alpha_a + \\alpha_b} |\\mathbf{A} - \\mathbf{B}|^2\\right)$$\n", + "\n", + "If this decay factor is below threshold for all primitive pairs, the shell pair contributes nothing." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Demonstrate shell-pair prescreening\n", + "print(\"Shell-Pair Prescreening:\")\n", + "print(\"-\" * 50)\n", + "\n", + "# Same center - always significant\n", + "shell_a = make_shell([0, 0, 0], 0, [1.0], [1.0])\n", + "shell_b = make_shell([0, 0, 0], 0, [0.5], [1.0])\n", + "print(f\"Same center (0,0,0): {shell_pair_significant(shell_a, shell_b)}\")\n", + "\n", + "# Close shells\n", + "shell_c = make_shell([1.0, 0, 0], 0, [1.0], [1.0])\n", + "print(f\"Close (1.0 Bohr apart): {shell_pair_significant(shell_a, shell_c)}\")\n", + "\n", + "# Far apart\n", + "shell_d = make_shell([50.0, 0, 0], 0, [1.0], [1.0])\n", + "print(f\"Far apart (50 Bohr): {shell_pair_significant(shell_a, shell_d)}\")\n", + "\n", + "# Very diffuse (small exponent) at distance\n", + "shell_e = make_shell([10.0, 0, 0], 0, [0.01], [1.0])\n", + "print(f\"Diffuse at 10 Bohr (a=0.01): {shell_pair_significant(shell_a, shell_e)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. High Angular Momentum Integrals\n", + "\n", + "The improved algorithm correctly handles **d-orbitals (L=2)** and **f-orbitals (L=3)** without numerical overflow." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "coord = np.array([0.0, 0.0, 0.0])\n", + "exps = np.array([0.5])\n", + "coeffs = np.array([[1.0]])\n", + "\n", + "# Angular momentum components\n", + "orbitals = {\n", + " 's (L=0)': (0, np.array([[0, 0, 0]])),\n", + " 'p (L=1)': (1, np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])),\n", + " 'd (L=2)': (2, np.array([[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0], [0, 1, 1], [0, 0, 2]])),\n", + " 'f (L=3)': (3, np.array([[3, 0, 0], [2, 1, 0], [2, 0, 1], [1, 2, 0], [1, 1, 1],\n", + " [1, 0, 2], [0, 3, 0], [0, 2, 1], [0, 1, 2], [0, 0, 3]]))\n", + "}\n", + "\n", + "print(\"High Angular Momentum Test:\")\n", + "print(\"-\" * 60)\n", + "print(f\"{'Orbital':>10} {'Shape':>20} {'Has NaN':>10} {'Has Inf':>10}\")\n", + "print(\"-\" * 60)\n", + "\n", + "for name, (angmom, comp) in orbitals.items():\n", + " result = compute_two_electron_integrals_os_hgp(\n", + " PointChargeIntegral.boys_func,\n", + " coord, angmom, comp, exps, coeffs,\n", + " coord, angmom, comp, exps, coeffs,\n", + " coord, angmom, comp, exps, coeffs,\n", + " coord, angmom, comp, exps, coeffs,\n", + " )\n", + " \n", + " shape_str = str(result.shape[:4])\n", + " has_nan = np.any(np.isnan(result))\n", + " has_inf = np.any(np.isinf(result))\n", + " \n", + " print(f\"{name:>10} {shape_str:>20} {str(has_nan):>10} {str(has_inf):>10}\")\n", + "\n", + "print(\"\\nAll integrals computed successfully without numerical issues!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Validation Against PySCF\n", + "\n", + "We validate our implementation against PySCF, a trusted quantum chemistry package.\n", + "Both the original and improved implementations should match PySCF exactly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " from pyscf import gto\n", + " HAS_PYSCF = True\n", + "except ImportError:\n", + " HAS_PYSCF = False\n", + " print(\"PySCF not installed. Install with: pip install pyscf\")\n", + "\n", + "if HAS_PYSCF:\n", + " # H2 with STO-3G basis\n", + " exps_h = np.array([3.42525091, 0.62391373, 0.16885540])\n", + " coeffs_h = np.array([[0.15432897], [0.53532814], [0.44463454]])\n", + " \n", + " basis_gbasis = [\n", + " GeneralizedContractionShell(0, np.array([0.0, 0.0, 0.0]), coeffs_h, exps_h, 'cartesian'),\n", + " GeneralizedContractionShell(0, np.array([1.4, 0.0, 0.0]), coeffs_h, exps_h, 'cartesian'),\n", + " ]\n", + " \n", + " mol = gto.M(\n", + " atom='H 0 0 0; H 1.4 0 0', unit='bohr',\n", + " basis={'H': gto.basis.parse(\n", + " 'H S\\n 3.42525091 0.15432897\\n 0.62391373 0.53532814\\n 0.16885540 0.44463454'\n", + " )}, verbose=0,\n", + " )\n", + " \n", + " # PySCF returns chemist notation (ij|kl)\n", + " eri_pyscf = mol.intor('int2e')\n", + " \n", + " # GBasis - both implementations in chemist notation\n", + " eri_original = electron_repulsion_integral(basis_gbasis, notation='chemist')\n", + " eri_improved = electron_repulsion_integral_improved(basis_gbasis, notation='chemist')\n", + " \n", + " print(\"H2/STO-3G Comparison (Chemist notation):\")\n", + " print(\"=\" * 55)\n", + " print(f\"{'Integral':>12} {'PySCF':>12} {'Original':>12} {'Improved':>12}\")\n", + " print(\"-\" * 55)\n", + " \n", + " for i in range(2):\n", + " for j in range(2):\n", + " label = f\"({i}{i}|{j}{j})\"\n", + " print(f\"{label:>12} {eri_pyscf[i,i,j,j]:>12.6f} {eri_original[i,i,j,j]:>12.6f} {eri_improved[i,i,j,j]:>12.6f}\")\n", + " \n", + " print(f\"\\nMax |Original - PySCF|: {np.max(np.abs(eri_original - eri_pyscf)):.2e}\")\n", + " print(f\"Max |Improved - PySCF|: {np.max(np.abs(eri_improved - eri_pyscf)):.2e}\")\n", + " print(f\"Max |Improved - Original|: {np.max(np.abs(eri_improved - eri_original)):.2e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 8. Summary\n\n### Key Features Implemented (Issue #221):\n\n| Week | Feature | Benefit |\n|------|---------|--------|\n| 1 | Boys Function | Efficient $F_m(T)$ via hyp1f1 for all orders |\n| 2 | VRR (Vertical Recursion) | Build angular momentum on center A |\n| 3 | ETR + Contraction | Transfer to center C + sum over primitives |\n| 4 | HRR (Horizontal Recursion) | Distribute to B and D (post-contraction per HGP) |\n| 5 | erf/erfc Boys Functions | Range-separated DFT support |\n| 6 | GBasis Integration | `ElectronRepulsionIntegralImproved` class |\n| 7 | Schwarz Screening | Skip negligible integrals (2-10x speedup) |\n| 8 | PySCF Validation | Verified against trusted reference |\n\n### When to Use Screening:\n- **Large molecules** with spatially separated atoms\n- **Extended basis sets** (cc-pVDZ, cc-pVTZ, etc.)\n- **Periodic systems** or **clusters**\n\n### When Screening Has Less Effect:\n- **Compact molecules** where all atoms are close\n- **Minimal basis sets** (STO-3G) with few primitives" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(\"\"\"\n", + "Thanks for completing this tutorial!\n", + "\n", + "For more information:\n", + "- Documentation: http://gbasis.qcdevs.org/\n", + "- GitHub: https://github.com/theochem/gbasis\n", + "- Issue #221: Two-electron integral improvements\n", + "\"\"\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/tests/test_boys_functions.py b/tests/test_boys_functions.py new file mode 100644 index 00000000..9870bd0d --- /dev/null +++ b/tests/test_boys_functions.py @@ -0,0 +1,559 @@ +"""Test gbasis.integrals.boys_functions module. + +Tests for standard Coulomb Boys function (Week 1) and +erf/erfc-attenuated potentials (Week 5). +""" + +import numpy as np +import pytest +from scipy.integrate import quad +from scipy.special import hyp1f1 + +from gbasis.integrals.boys_functions import ( + boys_function_all_orders, + boys_function_erf, + boys_function_erfc, + boys_function_mpmath, + boys_function_recursion, + boys_function_standard, + get_boys_function, +) + +try: + import mpmath # noqa: F401 + + HAS_MPMATH = True +except ImportError: + HAS_MPMATH = False + + +class TestBoysStandard: + """Tests for the standard Boys function (Coulomb potential).""" + + def test_boys_m0_zero(self): + """Test F_0(0) = 1.""" + result = boys_function_standard(np.array([0]), np.array([0.0])) + assert np.allclose(result, 1.0) + + def test_boys_m0_limit(self): + """Test F_0(T) -> sqrt(pi)/(2*sqrt(T)) for large T.""" + T = np.array([100.0]) + result = boys_function_standard(np.array([0]), T) + expected = np.sqrt(np.pi) / (2 * np.sqrt(T)) + assert np.allclose(result, expected, rtol=0.01) + + def test_boys_m1_zero(self): + """Test F_1(0) = 1/3.""" + result = boys_function_standard(np.array([1]), np.array([0.0])) + assert np.allclose(result, 1.0 / 3.0) + + def test_boys_general_zero(self): + """Test F_m(0) = 1/(2m+1).""" + for m in range(10): + result = boys_function_standard(np.array([m]), np.array([0.0])) + expected = 1.0 / (2 * m + 1) + assert np.allclose(result, expected), f"Failed for m={m}" + + def test_boys_consistency_with_hyp1f1(self): + """Test that our implementation matches scipy's hyp1f1.""" + orders = np.arange(5) + T_values = np.array([0.1, 0.5, 1.0, 2.0, 5.0]) + + for m in orders: + for T in T_values: + result = boys_function_standard(np.array([m]), np.array([T])) + expected = hyp1f1(m + 0.5, m + 1.5, -T) / (2 * m + 1) + assert np.allclose(result, expected), f"Failed for m={m}, T={T}" + + def test_boys_broadcasting(self): + """Test that Boys function handles broadcasting correctly (shape + values).""" + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + result = boys_function_standard(orders, T) + assert result.shape == (3, 3) + + # Verify each element matches individual scalar calls + for i, m in enumerate([0, 1, 2]): + for j, t in enumerate([0.5, 1.0, 2.0]): + expected = boys_function_standard(np.array([m]), np.array([t]))[0] + assert np.isclose( + result[i, j], expected + ), f"Broadcasting value mismatch at m={m}, T={t}" + + def test_boys_recurrence(self): + """Test the recurrence relation: (2m+1)*F_m(T) = 2T*F_{m+1}(T) + exp(-T).""" + T = np.array([1.0, 2.0, 3.0]) + for m in range(5): + fm = boys_function_standard(np.array([m]), T) + fm1 = boys_function_standard(np.array([m + 1]), T) + lhs = (2 * m + 1) * fm + rhs = 2 * T * fm1 + np.exp(-T) + assert np.allclose(lhs, rhs), f"Recurrence failed for m={m}" + + +class TestBoysAllOrders: + """Tests for the all-orders Boys function.""" + + def test_matches_standard_small_T(self): + """Test that all-orders function matches standard for small T.""" + m_max = 5 + T = np.array([0.1, 0.5, 1.0, 2.0]) + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + assert np.allclose(result[m], expected), f"Failed for m={m}" + + def test_matches_standard_large_T(self): + """Test that all-orders function matches standard for large T (>= 25).""" + m_max = 5 + T = np.array([30.0, 50.0, 100.0]) + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + np.testing.assert_allclose( + result[m], + expected.flatten(), + rtol=1e-8, + err_msg=f"All-orders doesn't match standard for m={m}, large T", + ) + + def test_output_shape(self): + """Test that all-orders function returns correct shape (m_max+1, *T.shape).""" + m_max = 4 + T = np.array([0.5, 1.0, 2.0, 5.0, 30.0]) + result = boys_function_all_orders(m_max, T) + assert result.shape == ( + m_max + 1, + T.shape[0], + ), f"Expected shape {(m_max + 1, T.shape[0])}, got {result.shape}" + + def test_small_T_series_branch(self): + """Exercise the small-T series path and compare with reference.""" + m_max = 6 + T = np.array([0.0, 1e-12, 1e-9]) + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + np.testing.assert_allclose( + result[m], expected, rtol=1e-12, atol=1e-14, err_msg=f"Small-T mismatch at m={m}" + ) + + def test_multidim_broadcasting(self): + """Check broadcasting for multidimensional T arrays.""" + m_max = 3 + T = np.array([[0.2, 1.0, 5.0], [0.4, 2.0, 10.0]])[:, None, :] + result = boys_function_all_orders(m_max, T) + + assert result.shape == (m_max + 1, *T.shape) + + # Verify values against scalar reference + for idx, t_val in np.ndenumerate(T): + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), np.array([t_val]))[0] + assert np.isclose(result[m][idx], expected), f"Mismatch at m={m}, idx={idx}" + + def test_recursion_identity(self): + """Verify recurrence holds on all-orders output.""" + m_max = 8 + T = np.array([0.3, 1.5, 4.0]) + all_vals = boys_function_all_orders(m_max, T) + + for m in range(m_max): + lhs = (2 * m + 1) * all_vals[m] + rhs = 2 * T * all_vals[m + 1] + np.exp(-T) + np.testing.assert_allclose( + lhs, rhs, rtol=1e-11, atol=1e-14, err_msg=f"Recurrence failed for m={m}" + ) + + def test_random_values_against_reference(self): + """Random spot-check across orders and T values.""" + rng = np.random.default_rng(0) + T = 10 ** rng.uniform(-10, 4, size=20) + m_max = 10 + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + np.testing.assert_allclose( + result[m], + expected, + rtol=1e-11, + atol=1e-14, + err_msg=f"Random check mismatch at m={m}", + ) + + +class TestGetBoysFunction: + """Tests for the get_boys_function factory.""" + + def test_coulomb_potential(self): + """Test that Coulomb potential returns standard Boys function.""" + boys = get_boys_function("coulomb") + result = boys(np.array([0]), np.array([1.0])) + expected = boys_function_standard(np.array([0]), np.array([1.0])) + assert np.allclose(result, expected) + + def test_aliases(self): + """Test that aliases work correctly.""" + for alias in ["coulomb", "standard", "1/r"]: + boys = get_boys_function(alias) + result = boys(np.array([0]), np.array([1.0])) + assert result is not None + + def test_unknown_potential(self): + """Test that unknown potential raises error.""" + with pytest.raises(ValueError): + get_boys_function("unknown_potential") + + def test_erf_potential(self): + """Test that erf potential returns correct function.""" + boys = get_boys_function("erf", omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + assert np.all(np.isfinite(result)) + + def test_erfc_potential(self): + """Test that erfc potential returns correct function.""" + boys = get_boys_function("erfc", omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + assert np.all(np.isfinite(result)) + + def test_erf_aliases(self): + """Test that erf aliases work correctly.""" + for alias in ["erf", "erf_coulomb"]: + boys = get_boys_function(alias, omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + + def test_erfc_aliases(self): + """Test that erfc aliases work correctly.""" + for alias in ["erfc", "erfc_coulomb"]: + boys = get_boys_function(alias, omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + + def test_erf_requires_omega(self): + """Test that erf potential requires omega parameter.""" + with pytest.raises(ValueError): + get_boys_function("erf") + + def test_erfc_requires_omega(self): + """Test that erfc potential requires omega parameter.""" + with pytest.raises(ValueError): + get_boys_function("erfc") + + def test_erf_requires_rho(self): + """Test that erf Boys function requires rho at call time.""" + boys = get_boys_function("erf", omega=0.5) + with pytest.raises(ValueError): + boys(np.array([0]), np.array([1.0])) + + def test_erfc_requires_rho(self): + """Test that erfc Boys function requires rho at call time.""" + boys = get_boys_function("erfc", omega=0.5) + with pytest.raises(ValueError): + boys(np.array([0]), np.array([1.0])) + + def test_case_insensitive(self): + """Test that potential matching is case-insensitive.""" + for alias in ["COULOMB", "Coulomb", "Standard", "1/R"]: + boys = get_boys_function(alias) + result = boys(np.array([0]), np.array([1.0])) + expected = boys_function_standard(np.array([0]), np.array([1.0])) + assert np.allclose(result, expected), f"Failed for '{alias}'" + + +class TestBoysNumericalIntegration: + """Compare Boys function with numerical integration.""" + + def test_numerical_integration_m0(self): + """Test F_0(T) against numerical integration.""" + + def boys_integrand(t, T, m): + return t ** (2 * m) * np.exp(-T * t**2) + + for T in [0.5, 1.0, 2.0, 5.0]: + result_analytic = boys_function_standard(np.array([0]), np.array([T]))[0] + result_numeric, _ = quad(boys_integrand, 0, 1, args=(T, 0)) + assert np.allclose(result_analytic, result_numeric, rtol=1e-6) + + def test_numerical_integration_m2(self): + """Test F_2(T) against numerical integration.""" + + def boys_integrand(t, T, m): + return t ** (2 * m) * np.exp(-T * t**2) + + for T in [0.5, 1.0, 2.0]: + result_analytic = boys_function_standard(np.array([2]), np.array([T]))[0] + result_numeric, _ = quad(boys_integrand, 0, 1, args=(T, 2)) + assert np.allclose(result_analytic, result_numeric, rtol=1e-6) + + +class TestBoysErf: + """Tests for the erf-attenuated Boys function.""" + + def test_erf_large_omega_recovers_coulomb(self): + """Test that erf with large omega approaches standard Coulomb. + + As omega -> infinity, erf(omega*r)/r -> 1/r. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erf = boys_function_erf(orders, T, rho, omega=1000.0) + result_std = boys_function_standard(orders, T) + + np.testing.assert_allclose( + result_erf, + result_std, + rtol=1e-5, + err_msg="erf with large omega should recover standard Coulomb", + ) + + def test_erf_small_omega_approaches_zero(self): + """Test that erf with small omega approaches zero. + + As omega -> 0, erf(omega*r)/r -> 0. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erf = boys_function_erf(orders, T, rho, omega=0.001) + result_std = boys_function_standard(orders, T) + # With omega=0.001, erf values should be much smaller than Coulomb + ratio = np.abs(result_erf) / (np.abs(result_std) + 1e-30) + assert np.all(ratio < 0.01), "erf with small omega should be << Coulomb" + + def test_erf_positive_values(self): + """Test that erf Boys function produces positive values.""" + orders = np.array([0, 1, 2, 3])[:, None] + T = np.array([0.1, 0.5, 1.0, 5.0])[None, :] + rho = np.ones_like(T) * 1.0 + + result = boys_function_erf(orders, T, rho, omega=0.5) + assert np.all(result >= 0), "erf Boys function should be non-negative" + + def test_erf_less_than_coulomb(self): + """Test that erf Boys < standard Coulomb (since erf(x) <= 1).""" + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erf = boys_function_erf(orders, T, rho, omega=0.5) + result_std = boys_function_standard(orders, T) + + assert np.all(result_erf <= result_std + 1e-15), "erf Boys should be <= standard Coulomb" + + def test_erf_scaling_formula(self): + """Test the erf scaling formula: scaling^(m+0.5) * F_m(scaling * T). + + Tests m=0,1,2 to ensure the exponent (m+0.5) is not hardcoded. + """ + T = np.array([1.0]) + rho = np.array([0.8]) + omega = 0.5 + scaling = omega**2 / (omega**2 + rho) + T_mod = scaling * T + + for m_val in [0, 1, 2]: + m = np.array([m_val]) + expected = scaling ** (m_val + 0.5) * boys_function_standard(m, T_mod) + result = boys_function_erf(m, T, rho, omega) + np.testing.assert_allclose( + result, expected, rtol=1e-12, err_msg=f"erf scaling formula failed for m={m_val}" + ) + + +class TestBoysErfc: + """Tests for the erfc-attenuated Boys function.""" + + def test_erfc_small_omega_recovers_coulomb(self): + """Test that erfc with small omega approaches standard Coulomb. + + As omega -> 0, erfc(omega*r)/r -> 1/r. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erfc = boys_function_erfc(orders, T, rho, omega=0.001) + result_std = boys_function_standard(orders, T) + + np.testing.assert_allclose( + result_erfc, + result_std, + rtol=0.01, + err_msg="erfc with small omega should recover standard Coulomb", + ) + + def test_erfc_large_omega_approaches_zero(self): + """Test that erfc with large omega approaches zero. + + As omega -> infinity, erfc(omega*r)/r -> 0. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erfc = boys_function_erfc(orders, T, rho, omega=1000.0) + assert np.all(np.abs(result_erfc) < 1e-6), "erfc with large omega should be near zero" + + def test_erf_plus_erfc_equals_coulomb(self): + """Test that erf + erfc = Coulomb for any omega. + + This is the fundamental identity: erf(x) + erfc(x) = 1, so + erf(omega*r)/r + erfc(omega*r)/r = 1/r. + """ + orders = np.array([0, 1, 2, 3])[:, None] + T = np.array([0.1, 0.5, 1.0, 2.0, 5.0])[None, :] + rho = np.ones_like(T) * 1.2 + + for omega in [0.1, 0.5, 1.0, 2.0]: + result_erf = boys_function_erf(orders, T, rho, omega) + result_erfc = boys_function_erfc(orders, T, rho, omega) + result_std = boys_function_standard(orders, T) + + np.testing.assert_allclose( + result_erf + result_erfc, + result_std, + rtol=1e-12, + err_msg=f"erf + erfc != Coulomb for omega={omega}", + ) + + def test_erfc_positive_values(self): + """Test that erfc Boys function produces positive values.""" + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.1, 0.5, 1.0, 5.0])[None, :] + rho = np.ones_like(T) * 1.0 + + result = boys_function_erfc(orders, T, rho, omega=0.5) + assert np.all(result >= -1e-15), "erfc Boys function should be non-negative" + + +class TestBoysRecursion: + """Tests for the downward recursion Boys function (Eq. 71).""" + + def test_recursion_matches_all_orders(self): + """Test that recursion matches hyp1f1-based all_orders for various T.""" + m_max = 10 + T = np.array([0.0, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 50.0, 100.0]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose( + result_rec, + result_ref, + rtol=1e-12, + atol=1e-15, + err_msg="Recursion doesn't match all_orders", + ) + + def test_recursion_T_zero(self): + """Test F_m(0) = 1/(2m+1) via recursion.""" + m_max = 15 + result = boys_function_recursion(m_max, np.array([0.0])) + for m in range(m_max + 1): + expected = 1.0 / (2 * m + 1) + np.testing.assert_allclose( + result[m], expected, rtol=1e-14, err_msg=f"F_{m}(0) incorrect" + ) + + def test_recursion_large_T(self): + """Test stability for large T values.""" + m_max = 10 + T = np.array([50.0, 100.0, 500.0]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose( + result_rec, result_ref, rtol=1e-10, err_msg="Recursion unstable for large T" + ) + + def test_recursion_high_order(self): + """Test with m_max = 20 (high angular momentum).""" + m_max = 20 + T = np.array([0.5, 2.0, 10.0]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose( + result_rec, result_ref, rtol=1e-11, err_msg="Recursion fails at high order" + ) + + def test_recursion_output_shape(self): + """Test output shape matches convention.""" + m_max = 5 + T = np.array([0.5, 1.0, 2.0]) + result = boys_function_recursion(m_max, T) + assert result.shape == (m_max + 1, 3) + + def test_recursion_multidim(self): + """Test with multidimensional T array.""" + m_max = 4 + T = np.array([[0.1, 1.0], [5.0, 10.0]]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + assert result_rec.shape == (m_max + 1, 2, 2) + np.testing.assert_allclose(result_rec, result_ref, rtol=1e-12) + + def test_recursion_recurrence_identity(self): + """Verify recurrence (2m+1)*F_m = 2T*F_{m+1} + exp(-T) holds.""" + m_max = 8 + T = np.array([0.3, 1.5, 4.0]) + all_vals = boys_function_recursion(m_max, T) + + for m in range(m_max): + lhs = (2 * m + 1) * all_vals[m] + rhs = 2 * T * all_vals[m + 1] + np.exp(-T) + np.testing.assert_allclose( + lhs, rhs, rtol=1e-11, atol=1e-14, err_msg=f"Recurrence failed for m={m}" + ) + + +@pytest.mark.skipif(not HAS_MPMATH, reason="mpmath not installed") +class TestBoysMpmath: + """Tests for the mpmath high-precision Boys function.""" + + def test_mpmath_matches_standard(self): + """Test that mpmath matches standard hyp1f1 for normal T values.""" + m_max = 5 + T = np.array([0.1, 0.5, 1.0, 5.0, 10.0]) + result_mp = boys_function_mpmath(m_max, T, dps=30) + result_std = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose( + result_mp, result_std, rtol=1e-14, err_msg="mpmath doesn't match standard for normal T" + ) + + def test_mpmath_output_shape(self): + """Test that mpmath returns correct shape.""" + m_max = 3 + T = np.array([[0.5, 1.0], [2.0, 5.0]]) + result = boys_function_mpmath(m_max, T, dps=20) + assert result.shape == (m_max + 1, 2, 2) + + def test_mpmath_T_zero(self): + """Test F_m(0) = 1/(2m+1) via mpmath.""" + m_max = 10 + result = boys_function_mpmath(m_max, np.array([0.0]), dps=50) + for m in range(m_max + 1): + expected = 1.0 / (2 * m + 1) + np.testing.assert_allclose(result[m], expected, rtol=1e-15) + + def test_mpmath_high_order(self): + """Test mpmath accuracy for high orders (m=15+).""" + m_max = 15 + T = np.array([1.0, 5.0]) + result_mp = boys_function_mpmath(m_max, T, dps=50) + result_std = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose( + result_mp, result_std, rtol=1e-14, err_msg="mpmath high-order mismatch" + ) diff --git a/tests/test_electron_repulsion.py b/tests/test_electron_repulsion.py index 0746b1db..c3d84525 100644 --- a/tests/test_electron_repulsion.py +++ b/tests/test_electron_repulsion.py @@ -1,217 +1,367 @@ -"""Test gbasis.integrals.electron_repulsion.""" +"""Test gbasis.integrals.electron_repulsion improved OS+HGP integration. + +Tests for Week 6: Integration of the improved OS+HGP algorithm into +the electron_repulsion module via ElectronRepulsionIntegralImproved class. +""" + +import numpy as np +import pytest +from utils import HortonContractions, find_datafile + from gbasis.contractions import GeneralizedContractionShell -from gbasis.integrals._two_elec_int import ( - _compute_two_elec_integrals, - _compute_two_elec_integrals_angmom_zero, -) from gbasis.integrals.electron_repulsion import ( - electron_repulsion_integral, ElectronRepulsionIntegral, + ElectronRepulsionIntegralImproved, + electron_repulsion_integral, + electron_repulsion_integral_improved, ) from gbasis.parsers import make_contractions, parse_nwchem -import numpy as np -import pytest -from utils import find_datafile, HortonContractions - - -def test_construct_array_contraction(): - """Test integrals.electron_repulsion.ElectronRepulsionIntegral.construct_array_contraction.""" - coord_one = np.array([0.5, 1, 1.5]) - cont_one = GeneralizedContractionShell( - 0, coord_one, np.array([1.0]), np.array([0.1]), "spherical" - ) - coord_two = np.array([1.5, 2, 3]) - cont_two = GeneralizedContractionShell( - 0, coord_two, np.array([3.0]), np.array([0.2]), "spherical" - ) - coord_three = np.array([2.5, 3, 4]) - cont_three = GeneralizedContractionShell( - 0, coord_three, np.array([3.0]), np.array([0.2]), "spherical" - ) - coord_four = np.array([3.5, 4, 5]) - cont_four = GeneralizedContractionShell( - 0, coord_four, np.array([3.0]), np.array([0.2]), "spherical" - ) - - with pytest.raises(TypeError): - ElectronRepulsionIntegral.construct_array_contraction(None, cont_two, cont_three, cont_four) - with pytest.raises(TypeError): - ElectronRepulsionIntegral.construct_array_contraction(cont_one, None, cont_three, cont_four) - with pytest.raises(TypeError): - ElectronRepulsionIntegral.construct_array_contraction(cont_one, cont_two, None, cont_four) - with pytest.raises(TypeError): - ElectronRepulsionIntegral.construct_array_contraction(cont_one, cont_two, cont_three, None) - - integrals = _compute_two_elec_integrals_angmom_zero( - ElectronRepulsionIntegral.boys_func, - cont_one.coord, - cont_one.exps, - cont_one.coeffs, - cont_two.coord, - cont_two.exps, - cont_two.coeffs, - cont_three.coord, - cont_three.exps, - cont_three.coeffs, - cont_four.coord, - cont_four.exps, - cont_four.coeffs, - ) - integrals = np.transpose(integrals, (4, 0, 5, 1, 6, 2, 7, 3)) - assert np.allclose( - integrals, - ElectronRepulsionIntegral.construct_array_contraction( - cont_one, cont_two, cont_three, cont_four - ), - ) - - cont_four.angmom = 1 - integrals = _compute_two_elec_integrals( - ElectronRepulsionIntegral.boys_func, - cont_one.coord, - cont_one.angmom, - cont_one.angmom_components_cart, - cont_one.exps, - cont_one.coeffs, - cont_two.coord, - cont_two.angmom, - cont_two.angmom_components_cart, - cont_two.exps, - cont_two.coeffs, - cont_three.coord, - cont_three.angmom, - cont_three.angmom_components_cart, - cont_three.exps, - cont_three.coeffs, - cont_four.coord, - cont_four.angmom, - cont_four.angmom_components_cart, - cont_four.exps, - cont_four.coeffs, - ) - integrals = np.transpose(integrals, (4, 0, 5, 1, 6, 2, 7, 3)) - assert np.allclose( - integrals, - ElectronRepulsionIntegral.construct_array_contraction( - cont_one, cont_two, cont_three, cont_four - ), - ) - - -def test_electron_repulsion_cartesian_horton_sto6g_bec(): - """Test electron_repulsion.electron_repulsion_cartesian against horton results. - - The test case is diatomic with Be and C separated by 1.0 angstroms with basis set STO-6G. Note - that ano-rcc was not used because it results in overflow in the _compute_two_electron_integrals. - - """ - basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) - coords = np.array([[0, 0, 0], [1.0, 0, 0]]) - basis = make_contractions(basis_dict, ["Be", "C"], coords, "cartesian") - basis = [HortonContractions(i.angmom, i.coord, i.coeffs, i.exps, i.coord_type) for i in basis] - - horton_elec_repulsion = np.load(find_datafile("data_horton_bec_cart_elec_repulsion.npy")) - assert np.allclose(horton_elec_repulsion, electron_repulsion_integral(basis)) - - -def test_electron_repulsion_cartesian_horton_custom_hhe(): - """Test electron_repulsion.electron_repulsion_cartesian against horton results. - - The test case is diatomic with H and He separated by 0.8 angstroms with basis set ANO-RCC - modified. The basis set was modified to remove large exponent components to avoid overflow and - some contractions for faster test. - - This test is also slow. - - """ - basis_dict = parse_nwchem(find_datafile("data_anorcc.nwchem")) - coords = np.array([[0, 0, 0], [0.8, 0, 0]]) - basis = make_contractions(basis_dict, ["H", "He"], coords, "cartesian") - basis = [ - HortonContractions(i.angmom, i.coord, i.coeffs[:, 0], i.exps, i.coord_type) - for i in basis[:8] - ] - basis[0] = HortonContractions( - basis[0].angmom, basis[0].coord, basis[0].coeffs[3:], basis[0].exps[3:], basis[0].coord_type - ) - basis[4] = HortonContractions( - basis[4].angmom, basis[4].coord, basis[4].coeffs[4:], basis[4].exps[4:], basis[4].coord_type - ) - basis.pop(3) - basis.pop(2) - - horton_elec_repulsion = np.load(find_datafile("data_horton_hhe_cart_elec_repulsion.npy")) - assert np.allclose(horton_elec_repulsion, electron_repulsion_integral(basis)) - - -def test_electron_repulsion_cartesian(): - """Test gbasis.integrals.electron_repulsion.electron_repulsion_cartesian.""" - basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) - basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "cartesian") - - erep_obj = ElectronRepulsionIntegral(basis) - assert np.allclose( - erep_obj.construct_array_cartesian(), - electron_repulsion_integral(basis, notation="chemist"), - ) - assert np.allclose( - np.einsum("ijkl->ikjl", erep_obj.construct_array_cartesian()), - electron_repulsion_integral(basis, notation="physicist"), - ) - with pytest.raises(ValueError): - electron_repulsion_integral(basis, notation="bad") - - -def test_electron_repulsion_spherical(): - """Test gbasis.integrals.electron_repulsion.electron_repulsion_spherical.""" - basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) - basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "spherical") - - erep_obj = ElectronRepulsionIntegral(basis) - assert np.allclose( - erep_obj.construct_array_spherical(), - electron_repulsion_integral(basis, notation="chemist"), - ) - assert np.allclose( - np.einsum("ijkl->ikjl", erep_obj.construct_array_spherical()), - electron_repulsion_integral(basis, notation="physicist"), - ) - with pytest.raises(ValueError): - electron_repulsion_integral(basis, notation="bad") - - -def test_electron_repulsion_mix(): - """Test gbasis.integrals.electron_repulsion.electron_repulsion_mix.""" - basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) - basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), ["spherical"] * 3) - - erep_obj = ElectronRepulsionIntegral(basis) - assert np.allclose( - erep_obj.construct_array_mix(["spherical"] * 3), - electron_repulsion_integral(basis, notation="chemist"), - ) - assert np.allclose( - np.einsum("ijkl->ikjl", erep_obj.construct_array_mix(["spherical"] * 3)), - electron_repulsion_integral(basis, notation="physicist"), - ) - with pytest.raises(ValueError): - electron_repulsion_integral(basis, notation="bad") - - -def test_electron_repulsion_lincomb(): - """Test gbasis.integrals.electron_repulsion.electron_repulsion_lincomb.""" - basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) - basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "spherical") - - erep_obj = ElectronRepulsionIntegral(basis) - transform = np.random.rand(3, 5) - assert np.allclose( - erep_obj.construct_array_lincomb(transform, ["spherical"]), - electron_repulsion_integral(basis, transform, notation="chemist"), - ) - assert np.allclose( - np.einsum("ijkl->ikjl", erep_obj.construct_array_lincomb(transform, ["spherical"])), - electron_repulsion_integral(basis, transform, notation="physicist"), - ) - with pytest.raises(ValueError): - electron_repulsion_integral(basis, transform, notation="bad") + + +class TestImprovedClassConstruction: + """Tests for ElectronRepulsionIntegralImproved class.""" + + def test_type_checks(self): + """Test that type checks raise TypeError for invalid inputs.""" + coord = np.array([0.0, 0.0, 0.0]) + cont = GeneralizedContractionShell(0, coord, np.array([1.0]), np.array([0.1]), "spherical") + with pytest.raises(TypeError): + ElectronRepulsionIntegralImproved.construct_array_contraction(None, cont, cont, cont) + with pytest.raises(TypeError): + ElectronRepulsionIntegralImproved.construct_array_contraction(cont, None, cont, cont) + with pytest.raises(TypeError): + ElectronRepulsionIntegralImproved.construct_array_contraction(cont, cont, None, cont) + with pytest.raises(TypeError): + ElectronRepulsionIntegralImproved.construct_array_contraction(cont, cont, cont, None) + + def test_ssss_contraction(self): + """Test (ss|ss) contraction with improved class.""" + coord_a = np.array([0.5, 1.0, 1.5]) + coord_b = np.array([1.5, 2.0, 3.0]) + coord_c = np.array([2.5, 3.0, 4.0]) + coord_d = np.array([3.5, 4.0, 5.0]) + + cont_a = GeneralizedContractionShell( + 0, coord_a, np.array([1.0]), np.array([0.1]), "spherical" + ) + cont_b = GeneralizedContractionShell( + 0, coord_b, np.array([3.0]), np.array([0.2]), "spherical" + ) + cont_c = GeneralizedContractionShell( + 0, coord_c, np.array([3.0]), np.array([0.2]), "spherical" + ) + cont_d = GeneralizedContractionShell( + 0, coord_d, np.array([3.0]), np.array([0.2]), "spherical" + ) + + result = ElectronRepulsionIntegralImproved.construct_array_contraction( + cont_a, cont_b, cont_c, cont_d + ) + assert result is not None + assert np.all(np.isfinite(result)) + + def test_output_shape_s_orbitals(self): + """Test output shape for s-orbital contractions.""" + coord = np.array([0.0, 0.0, 0.0]) + # coeffs shape (K, M) = (2, 1), exps shape (K,) = (2,) + cont = GeneralizedContractionShell( + 0, coord, np.array([[0.1], [0.2]]), np.array([1.0, 0.5]), "spherical" + ) + result = ElectronRepulsionIntegralImproved.construct_array_contraction( + cont, cont, cont, cont + ) + # Shape: (M_1, L_1, M_2, L_2, M_3, L_3, M_4, L_4) + # s-orbital: L=1, M depends on coeffs + assert result.shape[1] == 1 # L_cart for s-orbital + assert result.shape[3] == 1 + assert result.shape[5] == 1 + assert result.shape[7] == 1 + + +class TestImprovedMatchesOld: + """Test that improved implementation matches old implementation.""" + + def test_ssss_matches(self): + """Test (ss|ss) integrals match between old and improved.""" + coord_a = np.array([0.5, 1.0, 1.5]) + coord_b = np.array([1.5, 2.0, 3.0]) + coord_c = np.array([2.5, 3.0, 4.0]) + coord_d = np.array([3.5, 4.0, 5.0]) + + cont_a = GeneralizedContractionShell( + 0, coord_a, np.array([1.0]), np.array([0.1]), "spherical" + ) + cont_b = GeneralizedContractionShell( + 0, coord_b, np.array([3.0]), np.array([0.2]), "spherical" + ) + cont_c = GeneralizedContractionShell( + 0, coord_c, np.array([3.0]), np.array([0.2]), "spherical" + ) + cont_d = GeneralizedContractionShell( + 0, coord_d, np.array([3.0]), np.array([0.2]), "spherical" + ) + + result_old = ElectronRepulsionIntegral.construct_array_contraction( + cont_a, cont_b, cont_c, cont_d + ) + result_new = ElectronRepulsionIntegralImproved.construct_array_contraction( + cont_a, cont_b, cont_c, cont_d + ) + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="(ss|ss) integrals don't match between old and improved class", + ) + + def test_sssp_matches(self): + """Test (ss|sp) integrals match between old and improved.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + coord_d = np.array([1.0, 1.0, 0.0]) + + cont_s = GeneralizedContractionShell( + 0, coord_a, np.array([[1.0], [0.5]]), np.array([1.0, 0.5]), "spherical" + ) + cont_s2 = GeneralizedContractionShell( + 0, coord_b, np.array([1.0]), np.array([0.8]), "spherical" + ) + cont_s3 = GeneralizedContractionShell( + 0, coord_c, np.array([1.0]), np.array([1.2]), "spherical" + ) + cont_p = GeneralizedContractionShell( + 1, coord_d, np.array([1.0]), np.array([0.9]), "spherical" + ) + + result_old = ElectronRepulsionIntegral.construct_array_contraction( + cont_s, cont_s2, cont_s3, cont_p + ) + result_new = ElectronRepulsionIntegralImproved.construct_array_contraction( + cont_s, cont_s2, cont_s3, cont_p + ) + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="(ss|sp) integrals don't match between old and improved class", + ) + + def test_spsp_matches(self): + """Test (sp|sp) integrals match between old and improved.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.5, 0.0, 0.0]) + coord_c = np.array([0.0, 1.5, 0.0]) + coord_d = np.array([1.5, 1.5, 0.0]) + + cont_s = GeneralizedContractionShell( + 0, coord_a, np.array([[1.0], [0.5]]), np.array([1.0, 0.5]), "spherical" + ) + cont_p1 = GeneralizedContractionShell( + 1, coord_b, np.array([1.0]), np.array([0.8]), "spherical" + ) + cont_s2 = GeneralizedContractionShell( + 0, coord_c, np.array([1.0]), np.array([1.2]), "spherical" + ) + cont_p2 = GeneralizedContractionShell( + 1, coord_d, np.array([1.0]), np.array([0.9]), "spherical" + ) + + result_old = ElectronRepulsionIntegral.construct_array_contraction( + cont_s, cont_p1, cont_s2, cont_p2 + ) + result_new = ElectronRepulsionIntegralImproved.construct_array_contraction( + cont_s, cont_p1, cont_s2, cont_p2 + ) + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="(sp|sp) integrals don't match between old and improved class", + ) + + def test_pppp_matches(self): + """Test (pp|pp) integrals match between old and improved.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + coord_d = np.array([1.0, 1.0, 0.0]) + + cont_p1 = GeneralizedContractionShell( + 1, coord_a, np.array([1.0]), np.array([1.0]), "spherical" + ) + cont_p2 = GeneralizedContractionShell( + 1, coord_b, np.array([1.0]), np.array([0.8]), "spherical" + ) + cont_p3 = GeneralizedContractionShell( + 1, coord_c, np.array([1.0]), np.array([1.2]), "spherical" + ) + cont_p4 = GeneralizedContractionShell( + 1, coord_d, np.array([1.0]), np.array([0.9]), "spherical" + ) + + result_old = ElectronRepulsionIntegral.construct_array_contraction( + cont_p1, cont_p2, cont_p3, cont_p4 + ) + result_new = ElectronRepulsionIntegralImproved.construct_array_contraction( + cont_p1, cont_p2, cont_p3, cont_p4 + ) + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="(pp|pp) integrals don't match between old and improved class", + ) + + +class TestImprovedFullBasis: + """Test improved implementation with full basis sets.""" + + def test_sto6g_bec_cartesian_matches(self): + """Test improved matches old for Be-C with STO-6G basis (Cartesian).""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + coords = np.array([[0, 0, 0], [1.0, 0, 0]]) + basis = make_contractions(basis_dict, ["Be", "C"], coords, "cartesian") + basis = [ + HortonContractions(i.angmom, i.coord, i.coeffs, i.exps, i.coord_type) for i in basis + ] + + result_old = electron_repulsion_integral(basis, notation="chemist") + result_new = electron_repulsion_integral_improved(basis, notation="chemist") + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="STO-6G Be-C Cartesian integrals don't match", + ) + + def test_sto6g_bec_horton_reference(self): + """Test improved matches Horton reference for Be-C with STO-6G.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + coords = np.array([[0, 0, 0], [1.0, 0, 0]]) + basis = make_contractions(basis_dict, ["Be", "C"], coords, "cartesian") + basis = [ + HortonContractions(i.angmom, i.coord, i.coeffs, i.exps, i.coord_type) for i in basis + ] + + horton_ref = np.load(find_datafile("data_horton_bec_cart_elec_repulsion.npy")) + result_new = electron_repulsion_integral_improved(basis) + + np.testing.assert_allclose( + result_new, + horton_ref, + rtol=1e-10, + atol=1e-15, + err_msg="Improved integrals don't match Horton reference", + ) + + def test_sto6g_carbon_spherical_matches(self): + """Test improved matches old for Carbon with STO-6G (spherical).""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "spherical") + + result_old = electron_repulsion_integral(basis, notation="chemist") + result_new = electron_repulsion_integral_improved(basis, notation="chemist") + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="STO-6G Carbon spherical integrals don't match", + ) + + +class TestImprovedNotation: + """Test that notation conversions work correctly.""" + + def test_chemist_notation(self): + """Test that Chemists' notation works correctly.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "cartesian") + + erep_obj = ElectronRepulsionIntegralImproved(basis) + assert np.allclose( + erep_obj.construct_array_cartesian(), + electron_repulsion_integral_improved(basis, notation="chemist"), + ) + + def test_physicist_notation(self): + """Test that Physicists' notation works correctly.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "cartesian") + + erep_obj = ElectronRepulsionIntegralImproved(basis) + assert np.allclose( + np.einsum("ijkl->ikjl", erep_obj.construct_array_cartesian()), + electron_repulsion_integral_improved(basis, notation="physicist"), + ) + + def test_invalid_notation_raises(self): + """Test that invalid notation raises ValueError.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "cartesian") + + with pytest.raises(ValueError): + electron_repulsion_integral_improved(basis, notation="bad") + + def test_physicist_matches_old_physicist(self): + """Test that physicist notation matches between old and improved.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "cartesian") + + result_old = electron_repulsion_integral(basis, notation="physicist") + result_new = electron_repulsion_integral_improved(basis, notation="physicist") + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="Physicist notation doesn't match between old and improved", + ) + + +class TestImprovedSymmetries: + """Test that improved integrals satisfy expected symmetries.""" + + def test_chemist_symmetry_ijkl_jilk(self): + """Test (ij|kl) = (ji|lk) symmetry in Chemists' notation.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "spherical") + + integrals = electron_repulsion_integral_improved(basis, notation="chemist") + # (ij|kl) = (ji|lk) + np.testing.assert_allclose( + integrals, + np.transpose(integrals, (1, 0, 3, 2)), + rtol=1e-10, + err_msg="(ij|kl) != (ji|lk) symmetry violated", + ) + + def test_chemist_symmetry_ijkl_klij(self): + """Test (ij|kl) = (kl|ij) symmetry in Chemists' notation.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "spherical") + + integrals = electron_repulsion_integral_improved(basis, notation="chemist") + # (ij|kl) = (kl|ij) + np.testing.assert_allclose( + integrals, + np.transpose(integrals, (2, 3, 0, 1)), + rtol=1e-10, + err_msg="(ij|kl) != (kl|ij) symmetry violated", + ) + + def test_positive_diagonal(self): + """Test that diagonal integrals (ii|ii) are positive.""" + basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem")) + basis = make_contractions(basis_dict, ["C"], np.array([[0, 0, 0]]), "spherical") + + integrals = electron_repulsion_integral_improved(basis, notation="chemist") + n = integrals.shape[0] + for i in range(n): + assert integrals[i, i, i, i] > 0, f"Diagonal (ii|ii) not positive for i={i}" diff --git a/tests/test_libcint.py b/tests/test_libcint.py index 51c12667..1d1deddb 100644 --- a/tests/test_libcint.py +++ b/tests/test_libcint.py @@ -13,7 +13,10 @@ import gbasis from gbasis.integrals.angular_momentum import angular_momentum_integral -from gbasis.integrals.electron_repulsion import electron_repulsion_integral +from gbasis.integrals.electron_repulsion import ( + electron_repulsion_integral, + electron_repulsion_integral_improved, +) from gbasis.integrals.kinetic_energy import kinetic_energy_integral from gbasis.integrals.moment import moment_integral from gbasis.integrals.momentum import momentum_integral @@ -80,7 +83,7 @@ def test_integral(basis, atsyms, atcoords, coord_type, integral): against the GBasis Python integrals. """ - atol, rtol = 1e-4, 1e-4 + atol, rtol = 1e-6, 1e-6 atcoords = atcoords / 0.5291772083 @@ -125,7 +128,7 @@ def test_integral(basis, atsyms, atcoords, coord_type, integral): npt.assert_array_equal(lc_int.shape, (lc_basis.nbfn, lc_basis.nbfn, 3)) elif integral == "electron_repulsion": - py_int = electron_repulsion_integral(py_basis) + py_int = electron_repulsion_integral_improved(py_basis) npt.assert_array_equal( py_int.shape, (lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn) ) @@ -133,6 +136,8 @@ def test_integral(basis, atsyms, atcoords, coord_type, integral): npt.assert_array_equal( lc_int.shape, (lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn) ) + npt.assert_allclose(lc_int, py_int, atol=1e-4, rtol=1e-5) + return elif integral == "point_charge": charge_coords = np.asarray([[2.0, 2.0, 2.0], [-3.0, -3.0, -3.0], [-1.0, 2.0, -3.0]]) @@ -204,7 +209,7 @@ def test_integral_iodata(fname, elements, coord_type, integral, transform): from iodata import load_one from gbasis.integrals.libcint import ELEMENTS, LIBCINT, CBasis - atol, rtol = 1e-4, 1e-4 + atol, rtol = 1e-6, 1e-6 mol=load_one(find_datafile(fname)) py_basis=from_iodata(mol) diff --git a/tests/test_pyscf_comparison.py b/tests/test_pyscf_comparison.py new file mode 100644 index 00000000..fcded748 --- /dev/null +++ b/tests/test_pyscf_comparison.py @@ -0,0 +1,394 @@ +"""Direct comparison of GBasis ERIs with PySCF reference. + +This test validates that both the original and improved (OS+HGP) GBasis +implementations produce the same electron repulsion integrals as PySCF. + +Notation conventions: +- PySCF mol.intor('int2e') returns integrals in chemist's notation (ij|kl). +- GBasis electron_repulsion_integral() defaults to physicist's notation . +- GBasis electron_repulsion_integral(notation="chemist") returns chemist's (ij|kl). +- Conversion: = (ik|jl), so physicist = chemist.transpose(0, 2, 1, 3). + +We use manual basis definitions to ensure both codes use exactly the same +basis parameters, avoiding normalization conversion issues. +""" + +import numpy as np +import pytest + +from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals.electron_repulsion import ( + electron_repulsion_integral, + electron_repulsion_integral_improved, +) + +gto = pytest.importorskip("pyscf.gto", reason="pyscf not installed") + +# --------------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------------- + + +def make_h2_sto3g(): + """Create H2 with STO-3G basis (same parameters for GBasis and PySCF). + + STO-3G for H: 3 primitives contracted to 1 s-function. + Returns GBasis basis list and PySCF mol object. + """ + exps = np.array([3.42525091, 0.62391373, 0.16885540]) + coeffs = np.array([[0.15432897], [0.53532814], [0.44463454]]) + + coord1 = np.array([0.0, 0.0, 0.0]) + coord2 = np.array([1.4, 0.0, 0.0]) # 1.4 Bohr apart + + basis = [ + GeneralizedContractionShell(0, coord1, coeffs, exps, "cartesian"), + GeneralizedContractionShell(0, coord2, coeffs, exps, "cartesian"), + ] + + mol = gto.M( + atom="H 0 0 0; H 1.4 0 0", + unit="bohr", + basis={"H": gto.basis.parse(""" + H S + 3.42525091 0.15432897 + 0.62391373 0.53532814 + 0.16885540 0.44463454 + """)}, + verbose=0, + ) + return basis, mol + + +def make_h2_primitive(): + """Create H2 with a single primitive s-function per atom. + + Simplest possible case: no contraction complexity. + """ + exps = np.array([1.0]) + coeffs = np.array([[1.0]]) + + coord1 = np.array([0.0, 0.0, 0.0]) + coord2 = np.array([2.0, 0.0, 0.0]) + + basis = [ + GeneralizedContractionShell(0, coord1, coeffs, exps, "cartesian"), + GeneralizedContractionShell(0, coord2, coeffs, exps, "cartesian"), + ] + + mol = gto.M( + atom="H 0 0 0; H 2.0 0 0", + unit="bohr", + basis={"H": gto.basis.parse("H S\n 1.0 1.0")}, + verbose=0, + ) + return basis, mol + + +def make_h2_sp(): + """Create H2 with s + p basis (one s and one p per atom). + + This tests angular momentum handling, which is critical because + chemist/physicist notation differences matter for mixed angular momentum. + """ + s_exps = np.array([1.0]) + s_coeffs = np.array([[1.0]]) + p_exps = np.array([0.8]) + p_coeffs = np.array([[1.0]]) + + coord1 = np.array([0.0, 0.0, 0.0]) + coord2 = np.array([2.0, 0.0, 0.0]) + + basis = [ + GeneralizedContractionShell(0, coord1, s_coeffs, s_exps, "cartesian"), + GeneralizedContractionShell(1, coord1, p_coeffs, p_exps, "cartesian"), + GeneralizedContractionShell(0, coord2, s_coeffs, s_exps, "cartesian"), + GeneralizedContractionShell(1, coord2, p_coeffs, p_exps, "cartesian"), + ] + + mol = gto.M( + atom="H 0 0 0; H 2.0 0 0", + unit="bohr", + basis={"H": gto.basis.parse(""" + H S + 1.0 1.0 + H P + 0.8 1.0 + """)}, + verbose=0, + ) + return basis, mol + + +# --------------------------------------------------------------------------- +# Tests: Original implementation vs PySCF +# --------------------------------------------------------------------------- + + +class TestPySCFOriginal: + """Compare original GBasis ERIs with PySCF reference.""" + + def test_h2_sto3g_chemist(self): + """Test H2/STO-3G in chemist notation: GBasis (ij|kl) == PySCF (ij|kl).""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") # PySCF returns chemist (ij|kl) + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="H2/STO-3G chemist notation doesn't match PySCF", + ) + + def test_h2_sto3g_physicist(self): + """Test H2/STO-3G in physicist notation. + + GBasis physicist == PySCF chemist (ij|kl).transpose(0,2,1,3). + """ + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="physicist") + eri_pyscf = mol.intor("int2e") + eri_pyscf_physicist = eri_pyscf.transpose(0, 2, 1, 3) + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf_physicist, + rtol=1e-6, + atol=1e-10, + err_msg="H2/STO-3G physicist notation doesn't match PySCF", + ) + + def test_h2_primitive(self): + """Test single-primitive H2 (simplest possible case).""" + basis, mol = make_h2_primitive() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Single primitive H2 ERIs don't match PySCF", + ) + + def test_h2_sp_basis(self): + """Test H2 with s+p basis (tests angular momentum handling).""" + basis, mol = make_h2_sp() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="H2 s+p basis ERIs don't match PySCF", + ) + + def test_h2_specific_integrals(self): + """Test specific physically meaningful integrals for H2/STO-3G.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") + + # (00|00) - Coulomb integral, both electrons on atom 1 + assert np.isclose( + eri_gbasis[0, 0, 0, 0], eri_pyscf[0, 0, 0, 0], rtol=1e-6 + ), f"(00|00) mismatch: GBasis={eri_gbasis[0,0,0,0]:.10f}, PySCF={eri_pyscf[0,0,0,0]:.10f}" + + # (11|11) - Coulomb integral, both electrons on atom 2 + assert np.isclose( + eri_gbasis[1, 1, 1, 1], eri_pyscf[1, 1, 1, 1], rtol=1e-6 + ), f"(11|11) mismatch: GBasis={eri_gbasis[1,1,1,1]:.10f}, PySCF={eri_pyscf[1,1,1,1]:.10f}" + + # By symmetry, (00|00) == (11|11) for identical atoms + assert np.isclose( + eri_gbasis[0, 0, 0, 0], eri_gbasis[1, 1, 1, 1], rtol=1e-6 + ), "Identical atom diagonal integrals should be equal" + + # (01|01) - exchange-type integral + assert np.isclose( + eri_gbasis[0, 1, 0, 1], eri_pyscf[0, 1, 0, 1], rtol=1e-6 + ), "(01|01) mismatch" + + def test_h2_symmetries(self): + """Test that PySCF and GBasis agree on 8-fold ERI symmetries.""" + basis, _mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + + # Chemist notation symmetries: (ij|kl) = (ji|kl) = (ij|lk) = (kl|ij) + n = eri_gbasis.shape[0] + for i in range(n): + for j in range(n): + for k in range(n): + for m in range(n): + val = eri_gbasis[i, j, k, m] + assert np.isclose( + val, eri_gbasis[j, i, k, m], rtol=1e-10 + ), f"(ij|kl)!=(ji|kl) for ({i}{j}|{k}{m})" + assert np.isclose( + val, eri_gbasis[i, j, m, k], rtol=1e-10 + ), f"(ij|kl)!=(ij|lk) for ({i}{j}|{k}{m})" + assert np.isclose( + val, eri_gbasis[k, m, i, j], rtol=1e-10 + ), f"(ij|kl)!=(kl|ij) for ({i}{j}|{k}{m})" + + +# --------------------------------------------------------------------------- +# Tests: Improved (OS+HGP) implementation vs PySCF +# --------------------------------------------------------------------------- + + +class TestPySCFImproved: + """Compare improved OS+HGP GBasis ERIs with PySCF reference.""" + + def test_h2_sto3g_chemist(self): + """Test improved implementation H2/STO-3G in chemist notation.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Improved H2/STO-3G chemist notation doesn't match PySCF", + ) + + def test_h2_sto3g_physicist(self): + """Test improved implementation H2/STO-3G in physicist notation.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="physicist") + eri_pyscf = mol.intor("int2e") + eri_pyscf_physicist = eri_pyscf.transpose(0, 2, 1, 3) + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf_physicist, + rtol=1e-6, + atol=1e-10, + err_msg="Improved H2/STO-3G physicist notation doesn't match PySCF", + ) + + def test_h2_primitive(self): + """Test improved implementation with single primitive.""" + basis, mol = make_h2_primitive() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Improved single primitive H2 ERIs don't match PySCF", + ) + + def test_h2_sp_basis(self): + """Test improved implementation H2 with s+p basis.""" + basis, mol = make_h2_sp() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") + + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Improved H2 s+p basis ERIs don't match PySCF", + ) + + def test_improved_matches_original(self): + """Test that improved and original implementations agree on PySCF test case.""" + basis, _ = make_h2_sto3g() + + eri_old = electron_repulsion_integral(basis, notation="chemist") + eri_new = electron_repulsion_integral_improved(basis, notation="chemist") + + np.testing.assert_allclose( + eri_new, eri_old, rtol=1e-10, err_msg="Improved doesn't match original for H2/STO-3G" + ) + + def test_improved_matches_original_sp(self): + """Test that improved and original agree for s+p basis.""" + basis, _ = make_h2_sp() + + eri_old = electron_repulsion_integral(basis, notation="chemist") + eri_new = electron_repulsion_integral_improved(basis, notation="chemist") + + np.testing.assert_allclose( + eri_new, eri_old, rtol=1e-10, err_msg="Improved doesn't match original for H2 s+p" + ) + + +# --------------------------------------------------------------------------- +# Tests: Physical sanity checks +# --------------------------------------------------------------------------- + + +class TestPySCFPhysicalProperties: + """Verify physical properties of ERIs using PySCF as cross-check.""" + + def test_positive_diagonal(self): + """Test that Coulomb integrals (ii|ii) are positive.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor("int2e") + + n = eri_gbasis.shape[0] + for i in range(n): + assert eri_gbasis[i, i, i, i] > 0, f"GBasis (ii|ii) not positive for i={i}" + assert eri_pyscf[i, i, i, i] > 0, f"PySCF (ii|ii) not positive for i={i}" + + def test_coulomb_greater_than_exchange(self): + """Test that Coulomb integral >= exchange integral for H2. + + For same-atom shells: (00|00) >= (01|01) because the exchange + integral involves orbital overlap which reduces the value. + """ + basis, _mol = make_h2_sto3g() + + eri = electron_repulsion_integral_improved(basis, notation="chemist") + + # Coulomb (00|00) should be > exchange (01|01) + coulomb = eri[0, 0, 0, 0] + exchange = eri[0, 1, 0, 1] + assert coulomb > exchange, f"Coulomb ({coulomb:.6f}) should be > exchange ({exchange:.6f})" + + def test_eri_values_physically_reasonable(self): + """Test that ERI values are in a physically reasonable range. + + For H2 with STO-3G, typical values are 0.1-1.0 in atomic units. + """ + basis, _ = make_h2_sto3g() + + eri = electron_repulsion_integral_improved(basis, notation="chemist") + + # All values should be finite + assert np.all(np.isfinite(eri)), "ERIs contain NaN or Inf" + + # Diagonal integrals should be in reasonable range (0.1 - 2.0 a.u.) + for i in range(eri.shape[0]): + val = eri[i, i, i, i] + assert ( + 0.01 < val < 10.0 + ), f"Diagonal integral ({i}{i}|{i}{i}) = {val} outside reasonable range" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_schwarz_screening.py b/tests/test_schwarz_screening.py new file mode 100644 index 00000000..333f1552 --- /dev/null +++ b/tests/test_schwarz_screening.py @@ -0,0 +1,458 @@ +"""Benchmark tests for Schwarz screening and integral optimizations. + +Tests for Week 7: Schwarz screening (_screening.py) for efficient +two-electron integral computation. + +The optimizations tested: +1. Schwarz Screening - Skip negligible shell quartets using (ij|kl) <= Q_ij * Q_kl +2. Shell-pair Prescreening - Fast Gaussian decay check +3. SchwarzScreener class - Precomputed bounds with statistics +""" + +import numpy as np + +from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals._schwarz_screening import ( + SchwarzScreener, + compute_schwarz_bound_shell_pair, + compute_schwarz_bounds, + shell_pair_significant, +) +from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp +from gbasis.integrals.electron_repulsion import electron_repulsion_integral_improved +from gbasis.integrals.point_charge import PointChargeIntegral + + +def make_shell(coord, angmom, exps, coeffs): + """Create a GeneralizedContractionShell for testing.""" + return GeneralizedContractionShell( + angmom=angmom, + coord=np.array(coord), + coeffs=np.array(coeffs).reshape(-1, 1), + exps=np.array(exps), + coord_type="cartesian", + ) + + +def make_h_chain(n_atoms, spacing, include_p=False): + """Create a hydrogen chain basis set for testing. + + Parameters + ---------- + n_atoms : int + Number of hydrogen atoms in the chain. + spacing : float + Distance between atoms in Bohr. + include_p : bool + If True, include p-orbitals (cc-pVDZ-like basis). + + Returns + ------- + basis : list of GeneralizedContractionShell + Basis set for the hydrogen chain. + """ + # STO-3G exponents and coefficients + exps_s_sto3g = [3.42525091, 0.62391373, 0.16885540] + coeffs_s_sto3g = [0.15432897, 0.53532814, 0.44463454] + + # cc-pVDZ-like exponents + exps_s_dz = [13.01, 1.962, 0.4446, 0.122] + coeffs_s_dz = [0.0196, 0.1379, 0.4781, 0.5012] + exps_p = [0.727] + coeffs_p = [1.0] + + basis = [] + for i in range(n_atoms): + coord = [i * spacing, 0.0, 0.0] + if include_p: + basis.append(make_shell(coord, 0, exps_s_dz, coeffs_s_dz)) + basis.append(make_shell(coord, 1, exps_p, coeffs_p)) + else: + basis.append(make_shell(coord, 0, exps_s_sto3g, coeffs_s_sto3g)) + + return basis + + +class TestShellPairPrescreening: + """Tests for shell-pair significance screening.""" + + def test_same_center_always_significant(self): + """Test that shells on the same center are always significant.""" + shell = make_shell([0.0, 0.0, 0.0], 0, [1.0], [1.0]) + assert shell_pair_significant(shell, shell) is True + + def test_close_shells_significant(self): + """Test that close shells are significant.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [1.0], [1.0]) + shell_b = make_shell([1.0, 0.0, 0.0], 0, [1.0], [1.0]) + assert shell_pair_significant(shell_a, shell_b) is True + + def test_distant_shells_not_significant(self): + """Test that very distant shells with tight exponents are not significant.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [100.0], [1.0]) + shell_b = make_shell([100.0, 0.0, 0.0], 0, [100.0], [1.0]) + assert shell_pair_significant(shell_a, shell_b) is False + + def test_diffuse_shells_always_significant(self): + """Test that diffuse shells (small exponents) are significant even at distance.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [0.01], [1.0]) + shell_b = make_shell([10.0, 0.0, 0.0], 0, [0.01], [1.0]) + assert shell_pair_significant(shell_a, shell_b) is True + + def test_threshold_sensitivity(self): + """Test that threshold affects screening.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [5.0], [1.0]) + shell_b = make_shell([5.0, 0.0, 0.0], 0, [5.0], [1.0]) + + # With strict threshold, should be not significant + sig_strict = shell_pair_significant(shell_a, shell_b, threshold=1e-6) + # With loose threshold, should be significant + sig_loose = shell_pair_significant(shell_a, shell_b, threshold=1e-30) + + # Monotonicity: any pair significant under strict threshold must also be + # significant under loose threshold (smaller threshold = more permissive) + assert (not sig_strict) or sig_loose + + +class TestSchwarzBounds: + """Tests for Schwarz bound computation.""" + + def test_schwarz_bound_positive(self): + """Test that Schwarz bounds are non-negative.""" + shell = make_shell([0.0, 0.0, 0.0], 0, [1.0, 0.5], [0.6, 0.4]) + + boys_func = PointChargeIntegral.boys_func + bound = compute_schwarz_bound_shell_pair( + boys_func, shell, shell, compute_two_electron_integrals_os_hgp + ) + assert bound >= 0.0 + + def test_schwarz_bound_nonzero_for_same_shell(self): + """Test that Schwarz bound is nonzero for same shell.""" + shell = make_shell([0.0, 0.0, 0.0], 0, [1.0], [1.0]) + + boys_func = PointChargeIntegral.boys_func + bound = compute_schwarz_bound_shell_pair( + boys_func, shell, shell, compute_two_electron_integrals_os_hgp + ) + assert bound > 0.0 + + def test_schwarz_bounds_matrix_symmetric(self): + """Test that Schwarz bounds matrix is symmetric.""" + basis = make_h_chain(3, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + np.testing.assert_allclose( + bounds, bounds.T, rtol=1e-10, err_msg="Schwarz bounds matrix should be symmetric" + ) + + def test_schwarz_bounds_matrix_shape(self): + """Test that Schwarz bounds matrix has correct shape.""" + basis = make_h_chain(4, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + assert bounds.shape == (4, 4) + + def test_schwarz_bounds_diagonal_positive(self): + """Test that diagonal elements of bounds matrix are positive.""" + basis = make_h_chain(3, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + for i in range(len(basis)): + assert bounds[i, i] > 0, f"Diagonal bound[{i},{i}] should be positive" + + def test_schwarz_inequality_holds(self): + """Test that the Schwarz inequality |(ij|kl)| <= Q_ij * Q_kl holds.""" + basis = make_h_chain(3, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Compute actual integrals + integrals = electron_repulsion_integral_improved(basis, notation="chemist") + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + max_integral = np.max(np.abs(integrals[i, j, k, l_idx])) + schwarz_bound = bounds[i, j] * bounds[k, l_idx] + assert ( + max_integral <= schwarz_bound * (1 + 1e-6) + 1e-12 + ), f"Schwarz inequality violated for ({i}{j}|{k}{l_idx})" + + +class TestSchwarzScreener: + """Tests for the SchwarzScreener class.""" + + def test_screener_initialization(self): + """Test that screener initializes correctly.""" + basis = make_h_chain(3, spacing=2.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + assert screener.bounds.shape == (3, 3) + assert screener.n_screened == 0 + assert screener.n_computed == 0 + + def test_same_shell_significant(self): + """Test that same-shell quartets are always significant.""" + basis = make_h_chain(3, spacing=2.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + assert screener.is_significant(0, 0, 0, 0) is True + + def test_screening_statistics(self): + """Test that screening statistics are tracked correctly.""" + basis = make_h_chain(4, spacing=10.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Run through all quartets + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener.is_significant(i, j, k, l_idx) + + stats = screener.get_statistics() + assert stats["total"] == n**4 + assert stats["n_screened"] + stats["n_computed"] == n**4 + assert 0 <= stats["percent_screened"] <= 100 + + def test_extended_system_has_screening(self): + """Test that extended systems have significant screening.""" + basis = make_h_chain(4, spacing=15.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener.is_significant(i, j, k, l_idx) + + stats = screener.get_statistics() + # For well-separated atoms, many integrals should be screened + assert ( + stats["percent_screened"] > 30.0 + ), f"Expected >30% screening for extended system, got {stats['percent_screened']:.1f}%" + + def test_compact_system_no_screening(self): + """Test that compact systems have little screening.""" + basis = make_h_chain(3, spacing=1.4) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener.is_significant(i, j, k, l_idx) + + stats = screener.get_statistics() + # For compact molecules, most integrals should survive screening + assert ( + stats["percent_screened"] < 70.0 + ), f"Expected <70% screening for compact system, got {stats['percent_screened']:.1f}%" + + def test_reset_counters(self): + """Test that reset_counters works correctly.""" + basis = make_h_chain(3, spacing=2.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + screener.is_significant(0, 0, 0, 0) + screener.is_significant(0, 0, 1, 1) + assert screener.n_screened + screener.n_computed > 0 + + screener.reset_counters() + assert screener.n_screened == 0 + assert screener.n_computed == 0 + + def test_custom_threshold(self): + """Test that custom threshold affects screening.""" + basis = make_h_chain(4, spacing=5.0) + boys_func = PointChargeIntegral.boys_func + + # Strict threshold - more screening + screener_strict = SchwarzScreener( + basis, boys_func, compute_two_electron_integrals_os_hgp, threshold=1e-8 + ) + # Loose threshold - less screening + screener_loose = SchwarzScreener( + basis, boys_func, compute_two_electron_integrals_os_hgp, threshold=1e-16 + ) + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener_strict.is_significant(i, j, k, l_idx) + screener_loose.is_significant(i, j, k, l_idx) + + stats_strict = screener_strict.get_statistics() + stats_loose = screener_loose.get_statistics() + + # Strict threshold should screen at least as much as loose + assert ( + stats_strict["n_screened"] >= stats_loose["n_screened"] + ), "Strict threshold should screen more than loose threshold" + + +class TestScreeningCorrectness: + """Tests to verify that screened integrals match unscreened.""" + + def test_screened_vs_unscreened_compact(self): + """Test that screened computation gives same result for compact molecule.""" + basis = make_h_chain(3, spacing=1.4) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Compute full integrals + full_integrals = electron_repulsion_integral_improved(basis, notation="chemist") + + # Verify that significant quartets match + n = full_integrals.shape[0] + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + if not screener.is_significant(i, j, k, l_idx): + # Screened integrals should actually be negligible + assert ( + np.abs(full_integrals[i, j, k, l_idx]) < 1e-10 + ), f"Screened integral ({i}{j}|{k}{l_idx}) is not negligible" + + def test_screened_vs_unscreened_extended(self): + """Test that screened computation gives same result for extended molecule.""" + basis = make_h_chain(4, spacing=10.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Compute full integrals + full_integrals = electron_repulsion_integral_improved(basis, notation="chemist") + + n = full_integrals.shape[0] + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + if not screener.is_significant(i, j, k, l_idx): + assert ( + np.abs(full_integrals[i, j, k, l_idx]) < 1e-10 + ), f"Screened integral ({i}{j}|{k}{l_idx}) is not negligible" + + def test_screened_with_p_orbitals(self): + """Test screening correctness with p-orbitals.""" + basis = make_h_chain(2, spacing=8.0, include_p=True) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # For shell-level screening, verify bounds shape matches + assert screener.bounds.shape == (len(basis), len(basis)) + + # All diagonal bounds should be positive + for i in range(len(basis)): + assert screener.bounds[i, i] > 0 + + +class TestBenchmarkImprovedAlgorithm: + """Benchmark tests for the improved OS+HGP algorithm.""" + + def test_improved_algorithm_ssss(self): + """Test improved algorithm for (ss|ss) integrals.""" + coord = np.array([0.0, 0.0, 0.0]) + exps = np.array([1.0, 0.5]) + coeffs = np.array([[0.6], [0.4]]) + + angmom_s = 0 + angmom_comp_s = np.array([[0, 0, 0]]) + + boys_func = PointChargeIntegral.boys_func + + result = compute_two_electron_integrals_os_hgp( + boys_func, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + ) + + # Should be positive for (ss|ss) + assert result[0, 0, 0, 0, 0, 0, 0, 0] > 0 + + def test_improved_algorithm_pppp(self): + """Test improved algorithm for (pp|pp) integrals.""" + coord = np.array([0.0, 0.0, 0.0]) + exps = np.array([1.0]) + coeffs = np.array([[1.0]]) + + angmom_p = 1 + angmom_comp_p = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + + boys_func = PointChargeIntegral.boys_func + + result = compute_two_electron_integrals_os_hgp( + boys_func, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + ) + + assert result.shape == (3, 3, 3, 3, 1, 1, 1, 1) + assert result[0, 0, 0, 0, 0, 0, 0, 0] > 0 diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py new file mode 100644 index 00000000..5d9d9a1b --- /dev/null +++ b/tests/test_two_elec_int_improved.py @@ -0,0 +1,780 @@ +"""Test gbasis.integrals._two_elec_int_improved module. + +Tests for VRR (Week 2), ETR + contraction (Week 3), and HRR + full pipeline (Week 4). +""" + +import numpy as np +from scipy.special import hyp1f1 + +from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals._schwarz_screening import SchwarzScreener +from gbasis.integrals._two_elec_int import ( + _compute_two_elec_integrals, + _compute_two_elec_integrals_angmom_zero, +) +from gbasis.integrals._two_elec_int_improved import ( + _electron_transfer_recursion, + _get_factorial2_norm, + _optimized_contraction, + _vertical_recursion_relation, + compute_two_electron_integrals_os_hgp, +) +from gbasis.integrals.electron_repulsion import ( + ElectronRepulsionIntegralImproved, + electron_repulsion_integral_improved, +) +from gbasis.integrals.point_charge import PointChargeIntegral + + +class TestVerticalRecursion: + """Tests for the Vertical Recursion Relation (VRR).""" + + def test_vrr_base_case(self): + """Test that VRR preserves base case [00|00]^m.""" + m_max = 4 + n_prim = 2 + + # Create mock integrals_m (base case values) + integrals_m = np.random.rand(m_max, n_prim, n_prim, n_prim, n_prim) + + # Mock coordinates and exponents + rel_coord_a = np.random.rand(3, n_prim, n_prim, n_prim, n_prim) + coord_wac = np.random.rand(3, n_prim, n_prim, n_prim, n_prim) + harm_mean = np.random.rand(n_prim, n_prim, n_prim, n_prim) + 0.1 + exps_sum_one = np.random.rand(n_prim, n_prim, n_prim, n_prim) + 0.1 + + result = _vertical_recursion_relation( + integrals_m, m_max, rel_coord_a, coord_wac, harm_mean, exps_sum_one + ) + + # Base case should be preserved at [m, 0, 0, 0] + assert np.allclose(result[:, 0, 0, 0, ...], integrals_m) + + def test_vrr_output_shape(self): + """Test that VRR output has correct shape.""" + m_max = 5 + n_prim = 3 + + integrals_m = np.zeros((m_max, n_prim, n_prim, n_prim, n_prim)) + rel_coord_a = np.zeros((3, n_prim, n_prim, n_prim, n_prim)) + coord_wac = np.zeros((3, n_prim, n_prim, n_prim, n_prim)) + harm_mean = np.ones((n_prim, n_prim, n_prim, n_prim)) + exps_sum_one = np.ones((n_prim, n_prim, n_prim, n_prim)) + + result = _vertical_recursion_relation( + integrals_m, m_max, rel_coord_a, coord_wac, harm_mean, exps_sum_one + ) + + expected_shape = (m_max, m_max, m_max, m_max, n_prim, n_prim, n_prim, n_prim) + assert result.shape == expected_shape + + def test_vrr_p_orbital_manual(self): + """Test VRR first recursion step for p-orbital manually. + + For p-orbital, VRR computes: + [1,0|00]^0 = (P-A)_x * [00|00]^0 - (rho/zeta)*(Q-P)_x * [00|00]^1 + """ + m_max = 2 + integrals_m = np.array([[[[[1.0]]]], [[[[0.5]]]]]) + + PA_x, PA_y, PA_z = 0.3, 0.0, 0.0 + PQ_x, PQ_y, PQ_z = 0.2, 0.0, 0.0 + rho = 0.6 + zeta = 1.5 + + rel_coord_a = np.array([[[[[PA_x]]]], [[[[PA_y]]]], [[[[PA_z]]]]]) + coord_wac = np.array([[[[[PQ_x]]]], [[[[PQ_y]]]], [[[[PQ_z]]]]]) + harm_mean = np.array([[[[rho]]]]) + exps_sum_one = np.array([[[[zeta]]]]) + + result = _vertical_recursion_relation( + integrals_m, m_max, rel_coord_a, coord_wac, harm_mean, exps_sum_one + ) + + # Manual: [1,0,0|00]^0 = PA_x * [00|00]^0 - (rho/zeta)*PQ_x * [00|00]^1 + expected = PA_x * 1.0 - (rho / zeta) * PQ_x * 0.5 + assert np.allclose(result[0, 1, 0, 0, 0, 0, 0, 0], expected) + + def test_vrr_s_orbital_no_change(self): + """Test VRR with s-orbital where no recursion is needed.""" + m_max = 1 + integrals_m = np.array([[[[[3.14]]]]]) + + rel_coord_a = np.zeros((3, 1, 1, 1, 1)) + coord_wac = np.zeros((3, 1, 1, 1, 1)) + harm_mean = np.ones((1, 1, 1, 1)) + exps_sum_one = np.ones((1, 1, 1, 1)) + + result = _vertical_recursion_relation( + integrals_m, m_max, rel_coord_a, coord_wac, harm_mean, exps_sum_one + ) + + assert np.allclose(result[0, 0, 0, 0], 3.14) + + +class TestElectronTransferRecursion: + """Tests for the Electron Transfer Recursion (ETR).""" + + def test_etr_base_case(self): + """Test that ETR preserves base case (discards m, keeps a).""" + m_max = 4 + m_max_c = 3 + n_prim = 2 + + # Create mock VRR output + integrals_vert = np.random.rand(m_max, m_max, m_max, m_max, n_prim, n_prim, n_prim, n_prim) + + rel_coord_c = np.random.rand(3, n_prim, n_prim, n_prim, n_prim) + rel_coord_a = np.random.rand(3, n_prim, n_prim, n_prim, n_prim) + exps_sum_one = np.random.rand(n_prim, n_prim, n_prim, n_prim) + 0.1 + exps_sum_two = np.random.rand(n_prim, n_prim, n_prim, n_prim) + 0.1 + + result = _electron_transfer_recursion( + integrals_vert, m_max, m_max_c, rel_coord_c, rel_coord_a, exps_sum_one, exps_sum_two + ) + + # Base case: [0,0,0, a_x, a_y, a_z] should equal integrals_vert[0, a_x, a_y, a_z] + assert np.allclose(result[0, 0, 0, ...], integrals_vert[0, ...]) + + def test_etr_output_shape(self): + """Test that ETR output has correct shape.""" + m_max = 3 + m_max_c = 2 + n_prim = 2 + + integrals_vert = np.random.rand(m_max, m_max, m_max, m_max, n_prim, n_prim, n_prim, n_prim) + + rel_coord_c = np.random.rand(3, n_prim, n_prim, n_prim, n_prim) + rel_coord_a = np.random.rand(3, n_prim, n_prim, n_prim, n_prim) + exps_sum_one = np.random.rand(n_prim, n_prim, n_prim, n_prim) + 0.1 + exps_sum_two = np.random.rand(n_prim, n_prim, n_prim, n_prim) + 0.1 + + result = _electron_transfer_recursion( + integrals_vert, m_max, m_max_c, rel_coord_c, rel_coord_a, exps_sum_one, exps_sum_two + ) + + expected_shape = ( + m_max_c, + m_max_c, + m_max_c, + m_max, + m_max, + m_max, + n_prim, + n_prim, + n_prim, + n_prim, + ) + assert result.shape == expected_shape + + +class TestOptimizedContraction: + """Tests for the optimized primitive contraction.""" + + def test_output_shape(self): + """Test that contraction output has correct shape.""" + K, M = 2, 3 + integrals_etransf = np.random.rand(1, 1, 1, 1, 1, 1, K, K, K, K) + exps = np.random.rand(4, K) + 0.1 + coeffs = np.random.rand(4, K, M) + angmoms = np.array([0, 0, 0, 0]) + + result = _optimized_contraction(integrals_etransf, exps, coeffs, angmoms) + + expected_shape = (1, 1, 1, 1, 1, 1, M, M, M, M) + assert result.shape == expected_shape + + def test_accepts_tuples(self): + """Test that contraction accepts tuples as well as arrays.""" + K, M = 2, 2 + integrals_etransf = np.random.rand(1, 1, 1, 1, 1, 1, K, K, K, K) + exps = tuple(np.random.rand(K) + 0.1 for _ in range(4)) + coeffs = tuple(np.random.rand(K, M) for _ in range(4)) + angmoms = (0, 0, 0, 0) + + result = _optimized_contraction(integrals_etransf, exps, coeffs, angmoms) + + expected_shape = (1, 1, 1, 1, 1, 1, M, M, M, M) + assert result.shape == expected_shape + + +class TestFactorial2Norm: + """Tests for the factorial2 normalization helper.""" + + def test_s_orbital_norm(self): + """Test normalization for s-orbital (L=0).""" + s_key = ((0, 0, 0),) + norm = _get_factorial2_norm(s_key) + # (2*0-1)!! = (-1)!! = 1, so norm = 1/sqrt(1) = 1 + assert np.allclose(norm, 1.0) + + def test_p_orbital_norm(self): + """Test normalization for p-orbital (L=1).""" + p_key = ((1, 0, 0), (0, 1, 0), (0, 0, 1)) + norm = _get_factorial2_norm(p_key) + # Each component has one (2*1-1)!! = 1!! = 1 and two (2*0-1)!! = 1 + # So norm = 1/sqrt(1*1*1) = 1 for all + assert np.allclose(norm, 1.0) + + def test_caching(self): + """Test that factorial2 normalization is cached.""" + d_key = ((2, 0, 0), (1, 1, 0)) + norm1 = _get_factorial2_norm(d_key) + norm2 = _get_factorial2_norm(d_key) + assert np.allclose(norm1, norm2) + + +class TestImprovedVsOld: + """Compare improved OS+HGP implementation against old implementation.""" + + def test_ssss_matches_old(self): + """Test (ss|ss) integrals match old implementation. + + Note: The old _compute_two_elec_integrals has a known limitation with + all-zero angular momentum, so we compare against the specialized + _compute_two_elec_integrals_angmom_zero function instead. + """ + boys_func = PointChargeIntegral.boys_func + + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + coord_d = np.array([1.0, 1.0, 0.0]) + + angmom_comp = np.array([[0, 0, 0]]) + + exps_a = np.array([1.0, 0.5]) + exps_b = np.array([0.8]) + exps_c = np.array([1.2, 0.6]) + exps_d = np.array([0.9]) + + coeffs_a = np.array([[1.0], [0.5]]) + coeffs_b = np.array([[1.0]]) + coeffs_c = np.array([[1.0], [0.3]]) + coeffs_d = np.array([[1.0]]) + + # Old implementation (specialized for angmom zero) + result_old = _compute_two_elec_integrals_angmom_zero( + boys_func, + coord_a, + exps_a, + coeffs_a, + coord_b, + exps_b, + coeffs_b, + coord_c, + exps_c, + coeffs_c, + coord_d, + exps_d, + coeffs_d, + ) + + # New implementation + result_new = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 0, + angmom_comp, + exps_a, + coeffs_a, + coord_b, + 0, + angmom_comp, + exps_b, + coeffs_b, + coord_c, + 0, + angmom_comp, + exps_c, + coeffs_c, + coord_d, + 0, + angmom_comp, + exps_d, + coeffs_d, + ) + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="(ss|ss) integrals don't match between old and improved", + ) + + def test_spsp_matches_old(self): + """Test (sp|sp) integrals match old implementation.""" + boys_func = PointChargeIntegral.boys_func + + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.5, 0.0, 0.0]) + coord_c = np.array([0.0, 1.5, 0.0]) + coord_d = np.array([1.5, 1.5, 0.0]) + + exps_a = np.array([1.0, 0.5]) + exps_b = np.array([0.8]) + exps_c = np.array([1.2]) + exps_d = np.array([0.9]) + + coeffs_a = np.array([[1.0], [0.5]]) + coeffs_b = np.array([[1.0]]) + coeffs_c = np.array([[1.0]]) + coeffs_d = np.array([[1.0]]) + + # s-orbital components + s_comp = np.array([[0, 0, 0]]) + # p-orbital components + p_comp = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + + # Old implementation + result_old = _compute_two_elec_integrals( + boys_func, + coord_a, + 0, + s_comp, + exps_a, + coeffs_a, + coord_b, + 1, + p_comp, + exps_b, + coeffs_b, + coord_c, + 0, + s_comp, + exps_c, + coeffs_c, + coord_d, + 1, + p_comp, + exps_d, + coeffs_d, + ) + + # New implementation + result_new = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 0, + s_comp, + exps_a, + coeffs_a, + coord_b, + 1, + p_comp, + exps_b, + coeffs_b, + coord_c, + 0, + s_comp, + exps_c, + coeffs_c, + coord_d, + 1, + p_comp, + exps_d, + coeffs_d, + ) + + np.testing.assert_allclose( + result_new, + result_old, + rtol=1e-10, + err_msg="(sp|sp) integrals don't match between old and improved", + ) + + +class TestHighAngularMomentum: + """Test that high angular momentum integrals don't produce NaN/Inf.""" + + def test_dddd_no_overflow(self): + """Test (dd|dd) integrals produce finite values.""" + boys_func = PointChargeIntegral.boys_func + + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + coord_d = np.array([1.0, 1.0, 0.0]) + + exps = np.array([1.0]) + coeffs = np.array([[1.0]]) + + # d-orbital components (L=2): xx, xy, xz, yy, yz, zz + d_comp = np.array([[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0], [0, 1, 1], [0, 0, 2]]) + + result = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 2, + d_comp, + exps, + coeffs, + coord_b, + 2, + d_comp, + exps, + coeffs, + coord_c, + 2, + d_comp, + exps, + coeffs, + coord_d, + 2, + d_comp, + exps, + coeffs, + ) + + assert np.all(np.isfinite(result)), "d-orbital integrals contain NaN or Inf" + assert result.shape == (6, 6, 6, 6, 1, 1, 1, 1) + + +class TestPrimitiveScreening: + """Tests for primitive-level screening (Eq. 64).""" + + def test_screening_zero_threshold(self): + """With threshold=0, results match unscreened exactly.""" + + def boys_func(orders, weighted_dist, rho=None): + return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) + + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + exps = np.array([1.0, 0.5]) + coeffs = np.array([[1.0], [1.0]]) + s_comp = np.array([[0, 0, 0]]) + p_comp = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + + result_no_screen = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 1, + p_comp, + exps, + coeffs, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 1, + p_comp, + exps, + coeffs, + ) + + result_screened = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 1, + p_comp, + exps, + coeffs, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 1, + p_comp, + exps, + coeffs, + primitive_threshold=0.0, + ) + + np.testing.assert_array_equal(result_no_screen, result_screened) + + def test_screening_reasonable_threshold(self): + """With reasonable threshold, results match within tolerance.""" + + def boys_func(orders, weighted_dist, rho=None): + return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) + + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + exps = np.array([1.0, 0.5, 0.1]) + coeffs = np.array([[1.0], [1.0], [1.0]]) + s_comp = np.array([[0, 0, 0]]) + + result_no_screen = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 0, + s_comp, + exps, + coeffs, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 0, + s_comp, + exps, + coeffs, + ) + + result_screened = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 0, + s_comp, + exps, + coeffs, + coord_a, + 0, + s_comp, + exps, + coeffs, + coord_b, + 0, + s_comp, + exps, + coeffs, + primitive_threshold=1e-12, + ) + + np.testing.assert_allclose(result_no_screen, result_screened, atol=1e-10) + + +class TestContractionReordering: + """Tests for contraction reordering (l_a >= l_b, l_c >= l_d, l_a >= l_c).""" + + def test_sp_ps_reordering(self): + """Test that (sp|ps) gives correct results with bra/ket swapping.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + coord_d = np.array([1.0, 1.0, 0.0]) + + # Create shells: s (L=0) and p (L=1) + shell_s_a = GeneralizedContractionShell( + 0, coord_a, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_p_b = GeneralizedContractionShell( + 1, coord_b, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_p_c = GeneralizedContractionShell( + 1, coord_c, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s_d = GeneralizedContractionShell( + 0, coord_d, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + + # (sp|ps) triggers bra_swapped and ket_swapped + result = ElectronRepulsionIntegralImproved.construct_array_contraction( + shell_s_a, shell_p_b, shell_p_c, shell_s_d + ) + + # Verify result is finite and has correct shape + # Shape: (M_a, L_a, M_b, L_b, M_c, L_c, M_d, L_d) + assert result.shape == (1, 1, 1, 3, 1, 3, 1, 1) + assert np.all(np.isfinite(result)), "Reordered integrals contain NaN or Inf" + + # Compare with direct compute (no reordering) to verify correctness + boys_func = PointChargeIntegral.boys_func + + s_comp = np.array([[0, 0, 0]]) + p_comp = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + + direct = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 0, + s_comp, + np.array([1.0]), + np.array([[1.0]]), + coord_b, + 1, + p_comp, + np.array([1.0]), + np.array([[1.0]]), + coord_c, + 1, + p_comp, + np.array([1.0]), + np.array([[1.0]]), + coord_d, + 0, + s_comp, + np.array([1.0]), + np.array([[1.0]]), + ) + # direct shape: (L_a, L_b, L_c, L_d, M_a, M_b, M_c, M_d) + direct_transposed = np.transpose(direct, (4, 0, 5, 1, 6, 2, 7, 3)) + + np.testing.assert_allclose( + result, + direct_transposed, + rtol=1e-10, + err_msg="Reordered (sp|ps) doesn't match direct computation", + ) + + def test_pd_sp_braket_reordering(self): + """Test (pd|sp) triggers bra-ket swap as well.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + coord_d = np.array([1.0, 1.0, 0.0]) + + shell_p_a = GeneralizedContractionShell( + 1, coord_a, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_d_b = GeneralizedContractionShell( + 2, coord_b, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s_c = GeneralizedContractionShell( + 0, coord_c, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_p_d = GeneralizedContractionShell( + 1, coord_d, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + + # (pd|sp): bra_swapped (pp, no swap) + result = ElectronRepulsionIntegralImproved.construct_array_contraction( + shell_p_a, shell_d_b, shell_s_c, shell_p_d + ) + + # Shape: (M_a, L_p, M_b, L_d, M_c, L_s, M_d, L_p) + assert result.shape == (1, 3, 1, 6, 1, 1, 1, 3) + assert np.all(np.isfinite(result)), "Reordered integrals contain NaN or Inf" + + # Compare with direct (no reordering) + boys_func = PointChargeIntegral.boys_func + + s_comp = np.array([[0, 0, 0]]) + p_comp = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + d_comp = np.array([[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0], [0, 1, 1], [0, 0, 2]]) + + direct = compute_two_electron_integrals_os_hgp( + boys_func, + coord_a, + 1, + p_comp, + np.array([1.0]), + np.array([[1.0]]), + coord_b, + 2, + d_comp, + np.array([1.0]), + np.array([[1.0]]), + coord_c, + 0, + s_comp, + np.array([1.0]), + np.array([[1.0]]), + coord_d, + 1, + p_comp, + np.array([1.0]), + np.array([[1.0]]), + ) + direct_transposed = np.transpose(direct, (4, 0, 5, 1, 6, 2, 7, 3)) + + np.testing.assert_allclose( + result, + direct_transposed, + rtol=1e-10, + err_msg="Reordered (pd|sp) doesn't match direct computation", + ) + + +class TestSchwarzScreening: + """Tests for Schwarz screening integration.""" + + def test_schwarz_no_effect(self): + """With threshold=0, results match unscreened exactly.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + + shell_s = GeneralizedContractionShell( + 0, coord_a, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_p = GeneralizedContractionShell( + 1, coord_b, np.array([[1.0]]), np.array([0.8]), "cartesian" + ) + basis = [shell_s, shell_p] + + result_no_screen = electron_repulsion_integral_improved(basis, notation="chemist") + result_screened = electron_repulsion_integral_improved( + basis, notation="chemist", schwarz_threshold=0.0 + ) + + np.testing.assert_array_equal(result_no_screen, result_screened) + + def test_schwarz_screening_tight(self): + """With tight threshold, results match within tolerance.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + + shell_s1 = GeneralizedContractionShell( + 0, coord_a, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s2 = GeneralizedContractionShell( + 0, coord_b, np.array([[1.0]]), np.array([0.8]), "cartesian" + ) + shell_s3 = GeneralizedContractionShell( + 0, coord_c, np.array([[1.0]]), np.array([0.6]), "cartesian" + ) + basis = [shell_s1, shell_s2, shell_s3] + + result_no_screen = electron_repulsion_integral_improved(basis, notation="chemist") + result_screened = electron_repulsion_integral_improved( + basis, notation="chemist", schwarz_threshold=1e-12 + ) + + np.testing.assert_allclose(result_no_screen, result_screened, atol=1e-10) + + def test_schwarz_statistics(self): + """Verify screening statistics are tracked.""" + shell_s1 = GeneralizedContractionShell( + 0, np.array([0.0, 0.0, 0.0]), np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s2 = GeneralizedContractionShell( + 0, np.array([100.0, 0.0, 0.0]), np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s3 = GeneralizedContractionShell( + 0, np.array([0.0, 100.0, 0.0]), np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + basis = [shell_s1, shell_s2, shell_s3] + + screener = SchwarzScreener( + list(basis), + PointChargeIntegral.boys_func, + compute_two_electron_integrals_os_hgp, + threshold=1e-10, + ) + stats = screener.get_statistics() + # Far-apart shells should have some bounds below threshold + assert stats["total"] == 0 # No quartets checked yet via is_significant + + # Now run full computation with screening + result = electron_repulsion_integral_improved( + basis, notation="chemist", schwarz_threshold=1e-10 + ) + assert np.all(np.isfinite(result))