diff --git a/gbasis/integrals/_two_elec_int_improved.py b/gbasis/integrals/_two_elec_int_improved.py new file mode 100644 index 00000000..627fd2a4 --- /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, :, :, :, :]) * 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/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index 4643bbbe..68563f88 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -1,12 +1,11 @@ """Electron-electron repulsion integral.""" + +import numpy as np + from gbasis.base_four_symm import BaseFourIndexSymmetric 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._two_elec_int_improved import compute_two_electron_integrals_os_hgp from gbasis.integrals.point_charge import PointChargeIntegral -import numpy as np class ElectronRepulsionIntegral(BaseFourIndexSymmetric): @@ -156,46 +155,29 @@ def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four): raise TypeError("`cont_four` must be a `GeneralizedContractionShell` instance.") # TODO: we can probably swap the contractions to get the optimal time or memory usage - if cont_one.angmom == cont_two.angmom == cont_three.angmom == cont_four.angmom == 0: - integrals = _compute_two_elec_integrals_angmom_zero( - cls.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, - ) - else: - integrals = _compute_two_elec_integrals( - 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, - ) + 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, + ) integrals = np.transpose(integrals, (4, 0, 5, 1, 6, 2, 7, 3)) # TODO: if we swap the contractions, we need to unswap them here @@ -271,3 +253,179 @@ 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 = PointChargeIntegral.boys_func + + @classmethod + def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four): + 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.") + + # --- 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"): + 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. + + 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'") + + 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 + ) + elif all(ct == "cartesian" for ct in coord_type): + array = ElectronRepulsionIntegralImproved(basis).construct_array_cartesian() + elif all(ct == "spherical" for ct in coord_type): + array = ElectronRepulsionIntegralImproved(basis).construct_array_spherical() + else: + array = ElectronRepulsionIntegralImproved(basis).construct_array_mix(coord_type) + + if notation == "physicist": + array = np.transpose(array, (0, 2, 1, 3)) + return array 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_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py new file mode 100644 index 00000000..5974f117 --- /dev/null +++ b/tests/test_two_elec_int_improved.py @@ -0,0 +1,699 @@ +"""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._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 +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): + 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): + 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", + )