From 2ec8140282583a846f62c801080dce93850a1b52 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Tue, 3 Mar 2026 19:20:43 +0530 Subject: [PATCH 1/6] feat: implement ETR (Eq. 66) and primitive contraction for OS+HGP pipeline --- gbasis/integrals/_two_elec_int_improved.py | 386 +++++++++++++++++++++ 1 file changed, 386 insertions(+) create mode 100644 gbasis/integrals/_two_elec_int_improved.py diff --git a/gbasis/integrals/_two_elec_int_improved.py b/gbasis/integrals/_two_elec_int_improved.py new file mode 100644 index 00000000..eb088f9c --- /dev/null +++ b/gbasis/integrals/_two_elec_int_improved.py @@ -0,0 +1,386 @@ +"""Improved two-electron integrals using Obara-Saika + Head-Gordon-Pople recursions. + +This module implements VRR, ETR, and primitive contraction steps of the +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 +2. VRR: Build [a0|00]^m from [00|00]^m (Eq. 65) <-- Done (PR 2) +3. ETR: Build [a0|c0]^0 from [a0|00]^m (Eq. 66) <-- THIS PR +4. Contract primitives <-- THIS PR +5. HRR: Build [ab|cd] from [a0|c0] (Eq. 67) <-- Future + +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 numpy as np + +from gbasis.utils import factorial2 + +# Cache for factorial2 values to avoid repeated computation +_FACTORIAL2_CACHE = {} + + +def _get_factorial2_norm(angmom_components): + """Get cached factorial2 normalization for angular momentum components. + + Parameters + ---------- + angmom_components : np.ndarray(n, 3) + Angular momentum components. + + Returns + ------- + norm : np.ndarray(n,) + Normalization factors 1/sqrt(prod((2*l-1)!!)). + """ + key = tuple(map(tuple, angmom_components)) + if key not in _FACTORIAL2_CACHE: + _FACTORIAL2_CACHE[key] = 1.0 / np.sqrt( + np.prod(factorial2(2 * angmom_components - 1), axis=1) + ) + return _FACTORIAL2_CACHE[key] + + +def _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, +): + """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_a/b/c/d : np.ndarray + Primitive exponents. + coeffs_a/b/c/d : np.ndarray + Contraction coefficients. + angmom_a/b/c/d : int + Angular momenta. + + 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). + """ + # Precompute normalization constants (1D arrays) + norm_a = (2 * exps_a / np.pi) ** 0.75 * (4 * exps_a) ** (angmom_a / 2) + norm_b = (2 * exps_b / np.pi) ** 0.75 * (4 * exps_b) ** (angmom_b / 2) + norm_c = (2 * exps_c / np.pi) ** 0.75 * (4 * exps_c) ** (angmom_c / 2) + norm_d = (2 * exps_d / np.pi) ** 0.75 * (4 * exps_d) ** (angmom_d / 2) + + # Multiply coefficients by normalization (more efficient than per-element) + coeffs_a_norm = coeffs_a * norm_a[:, np.newaxis] + coeffs_b_norm = coeffs_b * norm_b[:, np.newaxis] + coeffs_c_norm = coeffs_c * norm_c[:, np.newaxis] + coeffs_d_norm = coeffs_d * norm_d[:, np.newaxis] + + # 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 From bf6e7273b058549ae0470e651dbbbd6d2a317056 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Tue, 3 Mar 2026 20:04:46 +0530 Subject: [PATCH 2/6] add test files for two elect int improved.py --- tests/test_two_elec_int_improved.py | 183 ++++++++++++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 tests/test_two_elec_int_improved.py diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py new file mode 100644 index 00000000..c870bc30 --- /dev/null +++ b/tests/test_two_elec_int_improved.py @@ -0,0 +1,183 @@ +"""Test gbasis.integrals._two_elec_int_improved module. + +Tests for VRR (Week 2) and ETR + contraction (Week 3). +""" + +import numpy as np +import pytest + +from gbasis.integrals._two_elec_int_improved import ( + _vertical_recursion_relation, + _electron_transfer_recursion, + _optimized_contraction, + _get_factorial2_norm, +) + + +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 TestFactorial2Norm: + """Tests for the factorial2 normalization helper.""" + + def test_s_orbital_norm(self): + """Test normalization for s-orbital (L=0).""" + s_components = np.array([[0, 0, 0]]) + norm = _get_factorial2_norm(s_components) + # (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_components = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + norm = _get_factorial2_norm(p_components) + # 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_components = np.array([[2, 0, 0], [1, 1, 0]]) + norm1 = _get_factorial2_norm(d_components) + norm2 = _get_factorial2_norm(d_components) + assert np.allclose(norm1, norm2) From d96fd22c3e190f6a5e88a1cdb0b7c429b0a88bfd Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Wed, 4 Mar 2026 18:30:56 +0530 Subject: [PATCH 3/6] style: apply black/ruff formatting fixes to test_two_elec_int_improved.py --- tests/test_two_elec_int_improved.py | 43 ++++++++++++++--------------- 1 file changed, 20 insertions(+), 23 deletions(-) diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py index c870bc30..bd4dff72 100644 --- a/tests/test_two_elec_int_improved.py +++ b/tests/test_two_elec_int_improved.py @@ -4,13 +4,11 @@ """ import numpy as np -import pytest from gbasis.integrals._two_elec_int_improved import ( - _vertical_recursion_relation, _electron_transfer_recursion, - _optimized_contraction, _get_factorial2_norm, + _vertical_recursion_relation, ) @@ -63,20 +61,15 @@ def test_vrr_p_orbital_manual(self): [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]]]]]) + 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]]]]]) + 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]]]]) @@ -115,9 +108,7 @@ def test_etr_base_case(self): 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 - ) + 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) @@ -125,8 +116,7 @@ def test_etr_base_case(self): 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 + 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] @@ -138,9 +128,7 @@ def test_etr_output_shape(self): 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 - ) + 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) @@ -148,12 +136,21 @@ def test_etr_output_shape(self): 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 + 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) + 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 From f0f57fa6854dbb497d355741fb7b78832929bf5a Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Thu, 5 Mar 2026 15:46:47 +0530 Subject: [PATCH 4/6] style: fix black formatting --- gbasis/integrals/_two_elec_int_improved.py | 57 +++++++++------------- tests/test_two_elec_int_improved.py | 14 +++--- 2 files changed, 29 insertions(+), 42 deletions(-) diff --git a/gbasis/integrals/_two_elec_int_improved.py b/gbasis/integrals/_two_elec_int_improved.py index eb088f9c..e9608dd8 100644 --- a/gbasis/integrals/_two_elec_int_improved.py +++ b/gbasis/integrals/_two_elec_int_improved.py @@ -16,68 +16,55 @@ - Ahlrichs, R. Phys. Chem. Chem. Phys. 2006, 8, 3072. """ +import functools + import numpy as np from gbasis.utils import factorial2 -# Cache for factorial2 values to avoid repeated computation -_FACTORIAL2_CACHE = {} - -def _get_factorial2_norm(angmom_components): +@functools.cache +def _get_factorial2_norm(angmom_key): """Get cached factorial2 normalization for angular momentum components. Parameters ---------- - angmom_components : np.ndarray(n, 3) - Angular momentum components. + 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)!!)). """ - key = tuple(map(tuple, angmom_components)) - if key not in _FACTORIAL2_CACHE: - _FACTORIAL2_CACHE[key] = 1.0 / np.sqrt( - np.prod(factorial2(2 * angmom_components - 1), axis=1) - ) - return _FACTORIAL2_CACHE[key] - - -def _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, -): + 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_a/b/c/d : np.ndarray - Primitive exponents. - coeffs_a/b/c/d : np.ndarray - Contraction coefficients. - angmom_a/b/c/d : int - Angular momenta. + exps : tuple of np.ndarray + Primitive exponents (exps_a, exps_b, exps_c, exps_d). + coeffs : tuple of np.ndarray + Contraction coefficients (coeffs_a, coeffs_b, coeffs_c, coeffs_d). + angmoms : tuple of int + Angular momenta (angmom_a, angmom_b, angmom_c, angmom_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). """ + exps_a, exps_b, exps_c, exps_d = exps + coeffs_a, coeffs_b, coeffs_c, coeffs_d = coeffs + angmom_a, angmom_b, angmom_c, angmom_d = angmoms + # Precompute normalization constants (1D arrays) norm_a = (2 * exps_a / np.pi) ** 0.75 * (4 * exps_a) ** (angmom_a / 2) norm_b = (2 * exps_b / np.pi) ** 0.75 * (4 * exps_b) ** (angmom_b / 2) diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py index bd4dff72..802a54f3 100644 --- a/tests/test_two_elec_int_improved.py +++ b/tests/test_two_elec_int_improved.py @@ -159,22 +159,22 @@ class TestFactorial2Norm: def test_s_orbital_norm(self): """Test normalization for s-orbital (L=0).""" - s_components = np.array([[0, 0, 0]]) - norm = _get_factorial2_norm(s_components) + 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_components = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) - norm = _get_factorial2_norm(p_components) + 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_components = np.array([[2, 0, 0], [1, 1, 0]]) - norm1 = _get_factorial2_norm(d_components) - norm2 = _get_factorial2_norm(d_components) + 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) From 2cea131fd98e960fe47075ec0598d4a85f27619b Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 6 Mar 2026 09:39:02 +0530 Subject: [PATCH 5/6] refactor: vectorize norm computation over all 4 centers using NumPy arrays --- gbasis/integrals/_two_elec_int_improved.py | 35 +++++++++------------- 1 file changed, 14 insertions(+), 21 deletions(-) diff --git a/gbasis/integrals/_two_elec_int_improved.py b/gbasis/integrals/_two_elec_int_improved.py index e9608dd8..2a23357d 100644 --- a/gbasis/integrals/_two_elec_int_improved.py +++ b/gbasis/integrals/_two_elec_int_improved.py @@ -49,33 +49,26 @@ def _optimized_contraction(integrals_etransf, exps, coeffs, angmoms): ---------- 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 : tuple of np.ndarray - Primitive exponents (exps_a, exps_b, exps_c, exps_d). - coeffs : tuple of np.ndarray - Contraction coefficients (coeffs_a, coeffs_b, coeffs_c, coeffs_d). - angmoms : tuple of int - Angular momenta (angmom_a, angmom_b, angmom_c, angmom_d). + 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). """ - exps_a, exps_b, exps_c, exps_d = exps - coeffs_a, coeffs_b, coeffs_c, coeffs_d = coeffs - angmom_a, angmom_b, angmom_c, angmom_d = angmoms - - # Precompute normalization constants (1D arrays) - norm_a = (2 * exps_a / np.pi) ** 0.75 * (4 * exps_a) ** (angmom_a / 2) - norm_b = (2 * exps_b / np.pi) ** 0.75 * (4 * exps_b) ** (angmom_b / 2) - norm_c = (2 * exps_c / np.pi) ** 0.75 * (4 * exps_c) ** (angmom_c / 2) - norm_d = (2 * exps_d / np.pi) ** 0.75 * (4 * exps_d) ** (angmom_d / 2) - - # Multiply coefficients by normalization (more efficient than per-element) - coeffs_a_norm = coeffs_a * norm_a[:, np.newaxis] - coeffs_b_norm = coeffs_b * norm_b[:, np.newaxis] - coeffs_c_norm = coeffs_c * norm_c[:, np.newaxis] - coeffs_d_norm = coeffs_d * norm_d[:, np.newaxis] + # Vectorized norm computation over all 4 centers + exps = np.array(exps) # shape (4, K) + angmoms = np.array(angmoms) # shape (4,) + coeffs = np.array(coeffs) # shape (4, K, M) + # shape (4, K) + norms = ((2 / np.pi) * exps) ** 0.75 * (4 * exps) ** (angmoms[:, np.newaxis] / 2) + coeffs_norm = coeffs * norms[:, :, np.newaxis] # shape (4, K, M) + 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) From 60ddebb85379cfe698c70a91988c24cf2eb111a5 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 6 Mar 2026 10:12:09 +0530 Subject: [PATCH 6/6] test: add tests for _optimized_contraction --- tests/test_two_elec_int_improved.py | 31 +++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py index 802a54f3..6ea7ddba 100644 --- a/tests/test_two_elec_int_improved.py +++ b/tests/test_two_elec_int_improved.py @@ -8,6 +8,7 @@ from gbasis.integrals._two_elec_int_improved import ( _electron_transfer_recursion, _get_factorial2_norm, + _optimized_contraction, _vertical_recursion_relation, ) @@ -154,6 +155,36 @@ def test_etr_output_shape(self): 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."""