From 2d3fe98026827c56552eed76ba502aa70a188d42 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Thu, 12 Feb 2026 15:25:22 +0530 Subject: [PATCH 01/29] feat: add Boys function APIs and validation tests --- .gitignore | 1 + gbasis/integrals/boys_functions.py | 132 ++++++++++++++++++++++++ tests/test_boys_functions.py | 157 +++++++++++++++++++++++++++++ 3 files changed, 290 insertions(+) create mode 100644 gbasis/integrals/boys_functions.py create mode 100644 tests/test_boys_functions.py diff --git a/.gitignore b/.gitignore index 7ffda729..aaf4569d 100644 --- a/.gitignore +++ b/.gitignore @@ -121,3 +121,4 @@ venv.bak/ # Libcint gbasis/integrals/lib/ gbasis/integrals/include/ +qc-dev-env/ diff --git a/gbasis/integrals/boys_functions.py b/gbasis/integrals/boys_functions.py new file mode 100644 index 00000000..eb08d257 --- /dev/null +++ b/gbasis/integrals/boys_functions.py @@ -0,0 +1,132 @@ +"""Boys functions for two-electron interaction potentials. + +This module implements the Boys function, which is the starting point for +computing two-electron integrals using the Obara-Saika recursion relations. + +The Boys function is defined as: + + F_m(T) = integral from 0 to 1 of t^(2m) * exp(-T*t^2) dt + +References: +- Helgaker, T.; Jorgensen, P.; Olsen, J. "Molecular Electronic-Structure Theory" + (eq. 9.8.39 for hyp1f1 representation) +- Ahlrichs, R. "A simple algebraic derivation of the Obara-Saika scheme for + general two-electron interaction potentials." + Phys. Chem. Chem. Phys., 2006, 8, 3072-3077. +""" + +import numpy as np +from scipy.special import hyp1f1 + + +def boys_function_standard(orders, weighted_dist): + r"""Compute standard Boys function for Coulomb potential (1/r12). + + The Coulombic Boys function is defined as: + + .. math:: + + F_m(T) = \int_0^1 t^{2m} e^{-T t^2} dt + + This can be expressed in terms of the Kummer confluent hypergeometric + function (hyp1f1) as shown in Helgaker (eq. 9.8.39). + + Parameters + ---------- + orders : np.ndarray + Differentiation order of the Boys function (m values). + Shape can be (M,) or (M, 1, 1, ...) for broadcasting. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + T = rho * |P - Q|^2 where rho = zeta*eta/(zeta+eta). + + Returns + ------- + boys_eval : np.ndarray + Boys function values F_m(T) with shape determined by broadcasting + of orders and weighted_dist. + + Notes + ----- + For the Coulomb potential g(r) = 1/r: + - G_0(rho, T) = (2*pi/rho) * F_0(T) + + This is the standard case used for electron repulsion integrals. + + There's some documented instability for hyp1f1, mainly for large values. + For typical quantum chemistry calculations (m < 20), the values are stable. + """ + return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) + + +def boys_function_all_orders(m_max, weighted_dist): + r"""Compute Boys function for all orders from 0 to m_max. + + Returns F_m(T) for m = 0, 1, ..., m_max simultaneously using + scipy's hyp1f1, which is numerically stable for all practical T + and m values encountered in quantum chemistry. + + Parameters + ---------- + m_max : int + Maximum order of the Boys function needed. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + + Returns + ------- + boys_all : np.ndarray + Boys function values for all orders from 0 to m_max. + Shape: (m_max + 1, *weighted_dist.shape) + """ + T = np.atleast_1d(weighted_dist) + result_shape = (m_max + 1, *T.shape) + boys_all = np.zeros(result_shape) + + # Use direct hyp1f1 formula for all T values. scipy's hyp1f1 is + # numerically stable for all practical T (tested up to T=100000) + # and m values encountered in quantum chemistry (m < 20). + for m in range(m_max + 1): + boys_all[m] = hyp1f1(m + 0.5, m + 1.5, -T) / (2 * m + 1) + + return boys_all + + +def get_boys_function(potential="coulomb", omega=None): + """Get the appropriate Boys function for a given potential type. + + Parameters + ---------- + potential : str + Type of two-electron potential. Options: + - "coulomb" or "standard" or "1/r": Standard 1/r12 Coulomb potential + omega : float, optional + Range-separation or damping parameter. Reserved for future use + with screened and damped potentials. + + Returns + ------- + boys_func : callable + Boys function with signature boys_func(orders, weighted_dist, rho=None). + + Raises + ------ + ValueError + If potential type is not recognized. + + Examples + -------- + >>> boys = get_boys_function("coulomb") + >>> F0 = boys(np.array([0]), np.array([1.0])) + """ + potential = potential.lower() + + if potential in ["coulomb", "standard", "1/r"]: + + def boys_func(orders, weighted_dist, rho=None): + return boys_function_standard(orders, weighted_dist) + + return boys_func + + else: + raise ValueError(f"Unknown potential type: {potential}. " f"Valid options: coulomb") diff --git a/tests/test_boys_functions.py b/tests/test_boys_functions.py new file mode 100644 index 00000000..810b5105 --- /dev/null +++ b/tests/test_boys_functions.py @@ -0,0 +1,157 @@ +"""Test gbasis.integrals.boys_functions module.""" + +import numpy as np +import pytest +from scipy.integrate import quad +from scipy.special import hyp1f1 + +from gbasis.integrals.boys_functions import ( + boys_function_all_orders, + boys_function_standard, + get_boys_function, +) + + +class TestBoysStandard: + """Tests for the standard Boys function (Coulomb potential).""" + + def test_boys_m0_zero(self): + """Test F_0(0) = 1.""" + result = boys_function_standard(np.array([0]), np.array([0.0])) + assert np.allclose(result, 1.0) + + def test_boys_m0_limit(self): + """Test F_0(T) -> sqrt(pi)/(2*sqrt(T)) for large T.""" + T = np.array([100.0]) + result = boys_function_standard(np.array([0]), T) + expected = np.sqrt(np.pi) / (2 * np.sqrt(T)) + assert np.allclose(result, expected, rtol=0.01) + + def test_boys_m1_zero(self): + """Test F_1(0) = 1/3.""" + result = boys_function_standard(np.array([1]), np.array([0.0])) + assert np.allclose(result, 1.0 / 3.0) + + def test_boys_general_zero(self): + """Test F_m(0) = 1/(2m+1).""" + for m in range(10): + result = boys_function_standard(np.array([m]), np.array([0.0])) + expected = 1.0 / (2 * m + 1) + assert np.allclose(result, expected), f"Failed for m={m}" + + def test_boys_consistency_with_hyp1f1(self): + """Test that our implementation matches scipy's hyp1f1.""" + orders = np.arange(5) + T_values = np.array([0.1, 0.5, 1.0, 2.0, 5.0]) + + for m in orders: + for T in T_values: + result = boys_function_standard(np.array([m]), np.array([T])) + expected = hyp1f1(m + 0.5, m + 1.5, -T) / (2 * m + 1) + assert np.allclose(result, expected), f"Failed for m={m}, T={T}" + + def test_boys_broadcasting(self): + """Test that Boys function handles broadcasting correctly (shape + values).""" + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + result = boys_function_standard(orders, T) + assert result.shape == (3, 3) + + # Verify each element matches individual scalar calls + for i, m in enumerate([0, 1, 2]): + for j, t in enumerate([0.5, 1.0, 2.0]): + expected = boys_function_standard(np.array([m]), np.array([t]))[0] + assert np.isclose(result[i, j], expected), \ + f"Broadcasting value mismatch at m={m}, T={t}" + + def test_boys_recurrence(self): + """Test the recurrence relation: (2m+1)*F_m(T) = 2T*F_{m+1}(T) + exp(-T).""" + T = np.array([1.0, 2.0, 3.0]) + for m in range(5): + fm = boys_function_standard(np.array([m]), T) + fm1 = boys_function_standard(np.array([m + 1]), T) + lhs = (2 * m + 1) * fm + rhs = 2 * T * fm1 + np.exp(-T) + assert np.allclose(lhs, rhs), f"Recurrence failed for m={m}" + + +class TestBoysAllOrders: + """Tests for the all-orders Boys function.""" + + def test_matches_standard_small_T(self): + """Test that all-orders function matches standard for small T.""" + m_max = 5 + T = np.array([0.1, 0.5, 1.0, 2.0]) + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + assert np.allclose(result[m], expected), f"Failed for m={m}" + + def test_matches_standard_large_T(self): + """Test that all-orders function matches standard for large T (>= 25).""" + m_max = 5 + T = np.array([30.0, 50.0, 100.0]) + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + np.testing.assert_allclose(result[m], expected.flatten(), rtol=1e-8, + err_msg=f"All-orders doesn't match standard for m={m}, large T") + + def test_output_shape(self): + """Test that all-orders function returns correct shape (m_max+1, *T.shape).""" + m_max = 4 + T = np.array([0.5, 1.0, 2.0, 5.0, 30.0]) + result = boys_function_all_orders(m_max, T) + assert result.shape == (m_max + 1, T.shape[0]), \ + f"Expected shape {(m_max + 1, T.shape[0])}, got {result.shape}" + + +class TestGetBoysFunction: + """Tests for the get_boys_function factory.""" + + def test_coulomb_potential(self): + """Test that Coulomb potential returns standard Boys function.""" + boys = get_boys_function("coulomb") + result = boys(np.array([0]), np.array([1.0])) + expected = boys_function_standard(np.array([0]), np.array([1.0])) + assert np.allclose(result, expected) + + def test_aliases(self): + """Test that aliases work correctly.""" + for alias in ["coulomb", "standard", "1/r"]: + boys = get_boys_function(alias) + result = boys(np.array([0]), np.array([1.0])) + assert result is not None + + def test_unknown_potential(self): + """Test that unknown potential raises error.""" + with pytest.raises(ValueError): + get_boys_function("unknown_potential") + + +class TestBoysNumericalIntegration: + """Compare Boys function with numerical integration.""" + + def test_numerical_integration_m0(self): + """Test F_0(T) against numerical integration.""" + + def boys_integrand(t, T, m): + return t ** (2 * m) * np.exp(-T * t ** 2) + + for T in [0.5, 1.0, 2.0, 5.0]: + result_analytic = boys_function_standard(np.array([0]), np.array([T]))[0] + result_numeric, _ = quad(boys_integrand, 0, 1, args=(T, 0)) + assert np.allclose(result_analytic, result_numeric, rtol=1e-6) + + def test_numerical_integration_m2(self): + """Test F_2(T) against numerical integration.""" + + def boys_integrand(t, T, m): + return t ** (2 * m) * np.exp(-T * t ** 2) + + for T in [0.5, 1.0, 2.0]: + result_analytic = boys_function_standard(np.array([2]), np.array([T]))[0] + result_numeric, _ = quad(boys_integrand, 0, 1, args=(T, 2)) + assert np.allclose(result_analytic, result_numeric, rtol=1e-6) From db5045e8fd5cd2ebf91832d239cd1e8539edbb24 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 13 Feb 2026 09:36:16 +0530 Subject: [PATCH 02/29] fix: move qc-dev-env/ from.gitignore to .git/info/exclude --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index aaf4569d..7ffda729 100644 --- a/.gitignore +++ b/.gitignore @@ -121,4 +121,3 @@ venv.bak/ # Libcint gbasis/integrals/lib/ gbasis/integrals/include/ -qc-dev-env/ From 4824a2a4c7274b9d6a6aabb4b12c48d7b98fb242 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 13 Feb 2026 16:50:51 +0530 Subject: [PATCH 03/29] refactor: vectorize boys_function_all_orders with NumPy broadcasting --- gbasis/integrals/boys_functions.py | 21 ++++++------ tests/test_boys_functions.py | 51 ++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/gbasis/integrals/boys_functions.py b/gbasis/integrals/boys_functions.py index eb08d257..884cf4ae 100644 --- a/gbasis/integrals/boys_functions.py +++ b/gbasis/integrals/boys_functions.py @@ -79,17 +79,20 @@ def boys_function_all_orders(m_max, weighted_dist): Boys function values for all orders from 0 to m_max. Shape: (m_max + 1, *weighted_dist.shape) """ - T = np.atleast_1d(weighted_dist) - result_shape = (m_max + 1, *T.shape) - boys_all = np.zeros(result_shape) + # Vectorize hyp1f1 across all orders with broadcasting to avoid + # per-order Python loops while keeping numerical behaviour identical + # to the standard expression. hyp1f1 is stable for the practical + # range m < 20 and T encountered here. - # Use direct hyp1f1 formula for all T values. scipy's hyp1f1 is - # numerically stable for all practical T (tested up to T=100000) - # and m values encountered in quantum chemistry (m < 20). - for m in range(m_max + 1): - boys_all[m] = hyp1f1(m + 0.5, m + 1.5, -T) / (2 * m + 1) + T = np.asarray(weighted_dist) + orders = np.arange(m_max + 1, dtype=np.result_type(T, np.float64)) - return boys_all + # Reshape orders to broadcast over T's dimensions: (m_max+1, 1, 1, ...) + orders_shape = (m_max + 1,) + (1,) * T.ndim + orders_b = orders.reshape(orders_shape) + + boys_vals = hyp1f1(orders_b + 0.5, orders_b + 1.5, -T) / (2 * orders_b + 1) + return boys_vals def get_boys_function(potential="coulomb", omega=None): diff --git a/tests/test_boys_functions.py b/tests/test_boys_functions.py index 810b5105..5ac205a7 100644 --- a/tests/test_boys_functions.py +++ b/tests/test_boys_functions.py @@ -107,6 +107,57 @@ def test_output_shape(self): assert result.shape == (m_max + 1, T.shape[0]), \ f"Expected shape {(m_max + 1, T.shape[0])}, got {result.shape}" + def test_small_T_series_branch(self): + """Exercise the small-T series path and compare with reference.""" + m_max = 6 + T = np.array([0.0, 1e-12, 1e-9]) + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + np.testing.assert_allclose( + result[m], expected, + rtol=1e-12, atol=1e-14, + err_msg=f"Small-T mismatch at m={m}") + + def test_multidim_broadcasting(self): + """Check broadcasting for multidimensional T arrays.""" + m_max = 3 + T = np.array([[0.2, 1.0, 5.0], [0.4, 2.0, 10.0]])[:, None, :] + result = boys_function_all_orders(m_max, T) + + assert result.shape == (m_max + 1, *T.shape) + + # Verify values against scalar reference + for idx, t_val in np.ndenumerate(T): + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), np.array([t_val]))[0] + assert np.isclose(result[m][idx], expected), f"Mismatch at m={m}, idx={idx}" + + def test_recursion_identity(self): + """Verify recurrence holds on all-orders output.""" + m_max = 8 + T = np.array([0.3, 1.5, 4.0]) + all_vals = boys_function_all_orders(m_max, T) + + for m in range(m_max): + lhs = (2 * m + 1) * all_vals[m] + rhs = 2 * T * all_vals[m + 1] + np.exp(-T) + np.testing.assert_allclose(lhs, rhs, rtol=1e-11, atol=1e-14, + err_msg=f"Recurrence failed for m={m}") + + def test_random_values_against_reference(self): + """Random spot-check across orders and T values.""" + rng = np.random.default_rng(0) + T = 10 ** rng.uniform(-10, 4, size=20) + m_max = 10 + result = boys_function_all_orders(m_max, T) + + for m in range(m_max + 1): + expected = boys_function_standard(np.array([m]), T) + np.testing.assert_allclose(result[m], expected, rtol=1e-11, atol=1e-14, + err_msg=f"Random check mismatch at m={m}") + class TestGetBoysFunction: """Tests for the get_boys_function factory.""" From 5278ebf097f97c89d15184515479e3e9a6d75a7d Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 27 Feb 2026 00:10:56 +0530 Subject: [PATCH 04/29] fix: add raise-from in mpmath ImportError (ruff B904) --- gbasis/integrals/boys_functions.py | 217 ++++++++++++++++++++++++++++- tests/test_boys_functions.py | 121 ++++++++++++++++ 2 files changed, 335 insertions(+), 3 deletions(-) diff --git a/gbasis/integrals/boys_functions.py b/gbasis/integrals/boys_functions.py index 884cf4ae..c3f04ed5 100644 --- a/gbasis/integrals/boys_functions.py +++ b/gbasis/integrals/boys_functions.py @@ -95,6 +95,187 @@ def boys_function_all_orders(m_max, weighted_dist): return boys_vals +def boys_function_recursion(m_max, weighted_dist): + r"""Compute Boys function for all orders using downward recursion. + + Computes F_m(T) for m = 0, 1, ..., m_max using the stable downward + recursion (Eq. 71 from algorithm notes): + + .. math:: + + F_m(T) = \frac{2T \cdot F_{m+1}(T) + e^{-T}}{2m + 1} + + The highest order F_{m_max}(T) is computed via hyp1f1 as the starting + value, then all lower orders are obtained by recursing downward. + This is more efficient than calling hyp1f1 for each order separately. + + Parameters + ---------- + m_max : int + Maximum order of the Boys function needed. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + + Returns + ------- + boys_all : np.ndarray + Boys function values for all orders from 0 to m_max. + Shape: (m_max + 1, *weighted_dist.shape) + """ + T = np.asarray(weighted_dist, dtype=np.float64) + result = np.empty((m_max + 1, *T.shape), dtype=np.float64) + + # Starting value: compute F_{m_max}(T) via hyp1f1 + result[m_max] = hyp1f1(m_max + 0.5, m_max + 1.5, -T) / (2 * m_max + 1) + + # Precompute exp(-T) once + exp_neg_T = np.exp(-T) + + # Downward recursion: F_m = (2T * F_{m+1} + exp(-T)) / (2m + 1) + for m in range(m_max - 1, -1, -1): + result[m] = (2.0 * T * result[m + 1] + exp_neg_T) / (2 * m + 1) + + return result + + +def boys_function_erf(orders, weighted_dist, rho, omega): + r"""Compute Boys function for erf-attenuated Coulomb potential. + + For the potential g(r) = erf(omega * r) / r, the modified Boys function is: + + .. math:: + + G_m(\rho, T) = \frac{2\pi}{\rho} \left(\frac{\omega^2}{\omega^2 + \rho}\right)^{m+1/2} + F_m\left(\frac{\omega^2 T}{\omega^2 + \rho}\right) + + This is used in range-separated DFT for the long-range part of the + electron-electron interaction. + + Parameters + ---------- + orders : np.ndarray + Differentiation orders (m values). + weighted_dist : np.ndarray + Weighted interatomic distances (T = rho * |P-Q|^2). + rho : np.ndarray + Harmonic mean of exponent sums: rho = zeta*eta/(zeta+eta). + omega : float + Range-separation parameter. + + Returns + ------- + boys_eval : np.ndarray + Modified Boys function values. + + Notes + ----- + As omega -> infinity, erf(omega*r)/r -> 1/r, recovering the standard Coulomb. + As omega -> 0, erf(omega*r)/r -> 0, the interaction vanishes. + """ + scaling = omega**2 / (omega**2 + rho) + T_modified = scaling * weighted_dist + return scaling ** (orders + 0.5) * boys_function_standard(orders, T_modified) + + +def boys_function_erfc(orders, weighted_dist, rho, omega): + r"""Compute Boys function for erfc-attenuated Coulomb potential. + + For the potential g(r) = erfc(omega * r) / r, which is the short-range + complement of the erf-attenuated potential: + + .. math:: + + G_m^{erfc} = G_m^{Coulomb} - G_m^{erf} + + This is used in range-separated DFT for the short-range part of the + electron-electron interaction. + + Parameters + ---------- + orders : np.ndarray + Differentiation orders (m values). + weighted_dist : np.ndarray + Weighted interatomic distances (T = rho * |P-Q|^2). + rho : np.ndarray + Harmonic mean of exponent sums: rho = zeta*eta/(zeta+eta). + omega : float + Range-separation parameter. + + Returns + ------- + boys_eval : np.ndarray + Modified Boys function values. + + Notes + ----- + erfc(omega*r)/r = 1/r - erf(omega*r)/r, so: + G_m^{erfc} = F_m(T) - (omega^2/(omega^2+rho))^(m+0.5) * F_m(omega^2*T/(omega^2+rho)) + + As omega -> infinity, erfc(omega*r)/r -> 0 (short-range vanishes). + As omega -> 0, erfc(omega*r)/r -> 1/r (full Coulomb). + """ + return boys_function_standard(orders, weighted_dist) - boys_function_erf( + orders, weighted_dist, rho, omega + ) + + +def boys_function_mpmath(m_max, weighted_dist, dps=50): + r"""Compute Boys function using mpmath for high-precision reference values. + + Uses mpmath.hyp1f1 with arbitrary precision arithmetic. Useful as a + reference for validating other implementations, or for edge cases where + scipy's hyp1f1 may lose precision (very high m or specific T ranges). + + Parameters + ---------- + m_max : int + Maximum order of the Boys function needed. + weighted_dist : np.ndarray + Weighted interatomic distances (T values). + dps : int + Decimal places of precision for mpmath (default: 50). + + Returns + ------- + boys_all : np.ndarray + Boys function values for all orders from 0 to m_max. + Shape: (m_max + 1, *weighted_dist.shape) + + Raises + ------ + ImportError + If mpmath is not installed. + + Notes + ----- + This function is SLOW compared to hyp1f1 or recursion-based methods. + It is intended only for validation and reference, not for production use. + Install mpmath via: pip install mpmath + """ + try: + import mpmath + except ImportError as err: + raise ImportError( + "mpmath is required for boys_function_mpmath. Install it with: pip install mpmath" + ) from err + + T = np.asarray(weighted_dist, dtype=np.float64) + result = np.empty((m_max + 1, *T.shape), dtype=np.float64) + + old_dps = mpmath.mp.dps + mpmath.mp.dps = dps + try: + for idx in np.ndindex(T.shape): + t_val = mpmath.mpf(float(T[idx])) + for m in range(m_max + 1): + val = mpmath.hyp1f1(m + 0.5, m + 1.5, -t_val) / (2 * m + 1) + result[(m, *idx)] = float(val) + finally: + mpmath.mp.dps = old_dps + + return result + + def get_boys_function(potential="coulomb", omega=None): """Get the appropriate Boys function for a given potential type. @@ -103,9 +284,11 @@ def get_boys_function(potential="coulomb", omega=None): potential : str Type of two-electron potential. Options: - "coulomb" or "standard" or "1/r": Standard 1/r12 Coulomb potential + - "erf" or "erf_coulomb": erf-attenuated (long-range) potential + - "erfc" or "erfc_coulomb": erfc-attenuated (short-range) potential omega : float, optional - Range-separation or damping parameter. Reserved for future use - with screened and damped potentials. + Range-separation parameter for erf/erfc potentials. + Required when potential is "erf" or "erfc". Returns ------- @@ -116,11 +299,15 @@ def get_boys_function(potential="coulomb", omega=None): ------ ValueError If potential type is not recognized. + If omega is not provided for erf/erfc potentials. Examples -------- >>> boys = get_boys_function("coulomb") >>> F0 = boys(np.array([0]), np.array([1.0])) + + >>> boys_erf = get_boys_function("erf", omega=0.4) + >>> F0_erf = boys_erf(np.array([0]), np.array([1.0]), rho=0.5) """ potential = potential.lower() @@ -131,5 +318,29 @@ def boys_func(orders, weighted_dist, rho=None): return boys_func + elif potential in ["erf", "erf_coulomb"]: + if omega is None: + raise ValueError("omega parameter is required for erf-attenuated potential") + + def boys_func(orders, weighted_dist, rho=None): + if rho is None: + raise ValueError("rho parameter is required for erf-attenuated potential") + return boys_function_erf(orders, weighted_dist, rho, omega) + + return boys_func + + elif potential in ["erfc", "erfc_coulomb"]: + if omega is None: + raise ValueError("omega parameter is required for erfc-attenuated potential") + + def boys_func(orders, weighted_dist, rho=None): + if rho is None: + raise ValueError("rho parameter is required for erfc-attenuated potential") + return boys_function_erfc(orders, weighted_dist, rho, omega) + + return boys_func + else: - raise ValueError(f"Unknown potential type: {potential}. " f"Valid options: coulomb") + raise ValueError( + f"Unknown potential type: {potential}. " f"Valid options: coulomb, erf, erfc" + ) diff --git a/tests/test_boys_functions.py b/tests/test_boys_functions.py index 5ac205a7..2741fb06 100644 --- a/tests/test_boys_functions.py +++ b/tests/test_boys_functions.py @@ -9,9 +9,18 @@ boys_function_all_orders, boys_function_standard, get_boys_function, + boys_function_recursion, + boys_function_mpmath, ) +try: + import mpmath # noqa: F401 + HAS_MPMATH = True +except ImportError: + HAS_MPMATH = False + + class TestBoysStandard: """Tests for the standard Boys function (Coulomb potential).""" @@ -206,3 +215,115 @@ def boys_integrand(t, T, m): result_analytic = boys_function_standard(np.array([2]), np.array([T]))[0] result_numeric, _ = quad(boys_integrand, 0, 1, args=(T, 2)) assert np.allclose(result_analytic, result_numeric, rtol=1e-6) + +class TestBoysRecursion: + """Tests for the downward recursion Boys function (Eq. 71).""" + + def test_recursion_matches_all_orders(self): + """Test that recursion matches hyp1f1-based all_orders for various T.""" + m_max = 10 + T = np.array([0.0, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 50.0, 100.0]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose(result_rec, result_ref, rtol=1e-12, atol=1e-15, + err_msg="Recursion doesn't match all_orders") + + def test_recursion_T_zero(self): + """Test F_m(0) = 1/(2m+1) via recursion.""" + m_max = 15 + result = boys_function_recursion(m_max, np.array([0.0])) + for m in range(m_max + 1): + expected = 1.0 / (2 * m + 1) + np.testing.assert_allclose(result[m], expected, rtol=1e-14, + err_msg=f"F_{m}(0) incorrect") + + def test_recursion_large_T(self): + """Test stability for large T values.""" + m_max = 10 + T = np.array([50.0, 100.0, 500.0]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose(result_rec, result_ref, rtol=1e-10, + err_msg="Recursion unstable for large T") + + def test_recursion_high_order(self): + """Test with m_max = 20 (high angular momentum).""" + m_max = 20 + T = np.array([0.5, 2.0, 10.0]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose(result_rec, result_ref, rtol=1e-11, + err_msg="Recursion fails at high order") + + def test_recursion_output_shape(self): + """Test output shape matches convention.""" + m_max = 5 + T = np.array([0.5, 1.0, 2.0]) + result = boys_function_recursion(m_max, T) + assert result.shape == (m_max + 1, 3) + + def test_recursion_multidim(self): + """Test with multidimensional T array.""" + m_max = 4 + T = np.array([[0.1, 1.0], [5.0, 10.0]]) + result_rec = boys_function_recursion(m_max, T) + result_ref = boys_function_all_orders(m_max, T) + + assert result_rec.shape == (m_max + 1, 2, 2) + np.testing.assert_allclose(result_rec, result_ref, rtol=1e-12) + + def test_recursion_recurrence_identity(self): + """Verify recurrence (2m+1)*F_m = 2T*F_{m+1} + exp(-T) holds.""" + m_max = 8 + T = np.array([0.3, 1.5, 4.0]) + all_vals = boys_function_recursion(m_max, T) + + for m in range(m_max): + lhs = (2 * m + 1) * all_vals[m] + rhs = 2 * T * all_vals[m + 1] + np.exp(-T) + np.testing.assert_allclose(lhs, rhs, rtol=1e-11, atol=1e-14, + err_msg=f"Recurrence failed for m={m}") + + +@pytest.mark.skipif(not HAS_MPMATH, reason="mpmath not installed") +class TestBoysMpmath: + """Tests for the mpmath high-precision Boys function.""" + + def test_mpmath_matches_standard(self): + """Test that mpmath matches standard hyp1f1 for normal T values.""" + m_max = 5 + T = np.array([0.1, 0.5, 1.0, 5.0, 10.0]) + result_mp = boys_function_mpmath(m_max, T, dps=30) + result_std = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose(result_mp, result_std, rtol=1e-14, + err_msg="mpmath doesn't match standard for normal T") + + def test_mpmath_output_shape(self): + """Test that mpmath returns correct shape.""" + m_max = 3 + T = np.array([[0.5, 1.0], [2.0, 5.0]]) + result = boys_function_mpmath(m_max, T, dps=20) + assert result.shape == (m_max + 1, 2, 2) + + def test_mpmath_T_zero(self): + """Test F_m(0) = 1/(2m+1) via mpmath.""" + m_max = 10 + result = boys_function_mpmath(m_max, np.array([0.0]), dps=50) + for m in range(m_max + 1): + expected = 1.0 / (2 * m + 1) + np.testing.assert_allclose(result[m], expected, rtol=1e-15) + + def test_mpmath_high_order(self): + """Test mpmath accuracy for high orders (m=15+).""" + m_max = 15 + T = np.array([1.0, 5.0]) + result_mp = boys_function_mpmath(m_max, T, dps=50) + result_std = boys_function_all_orders(m_max, T) + + np.testing.assert_allclose(result_mp, result_std, rtol=1e-14, + err_msg="mpmath high-order mismatch") + From c19cb89883201ff9fad346563fb967d7eddd04ad Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 27 Feb 2026 01:14:26 +0530 Subject: [PATCH 05/29] feat: implement VRR (Eq. 65) for OS+HGP two-electron integral pipeline --- gbasis/integrals/_two_elec_int_improved.py | 131 +++++++++++++++++++++ tests/test_two_elec_int_improved.py | 102 ++++++++++++++++ 2 files changed, 233 insertions(+) create mode 100644 gbasis/integrals/_two_elec_int_improved.py create mode 100644 tests/test_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..a78b74f8 --- /dev/null +++ b/gbasis/integrals/_two_elec_int_improved.py @@ -0,0 +1,131 @@ +"""Improved two-electron integrals using Obara-Saika + Head-Gordon-Pople recursions. + +This module implements the Vertical Recursion Relation (VRR) for building +angular momentum on center A, as the first step of the OS+HGP algorithm. + +Algorithm overview (full pipeline, to be completed in future PRs): +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) <-- THIS PR +3. ETR: Build [a0|c0]^0 from [a0|00]^m (Eq. 66) <-- Future +4. Contract primitives <-- Future +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 + + +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 diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py new file mode 100644 index 00000000..da68824a --- /dev/null +++ b/tests/test_two_elec_int_improved.py @@ -0,0 +1,102 @@ +"""Test gbasis.integrals._two_elec_int_improved module. + +Week 2: Tests for VRR (Vertical Recursion Relation). +""" + +import numpy as np +import pytest + +from gbasis.integrals._two_elec_int_improved import ( + _vertical_recursion_relation, +) + + +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) From 2ec8140282583a846f62c801080dce93850a1b52 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Tue, 3 Mar 2026 19:20:43 +0530 Subject: [PATCH 06/29] 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 6e40ec77bca33b50aca0d3ec1efb52ae817d4f22 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Tue, 3 Mar 2026 19:53:06 +0530 Subject: [PATCH 07/29] feat: implement HRR (Eq. 67) and complete OS+HGP pipeline --- gbasis/integrals/_two_elec_int_improved.py | 757 +++++++++++++++++++++ tests/test_two_elec_int_improved.py | 392 +++++++++++ 2 files changed, 1149 insertions(+) create mode 100644 gbasis/integrals/_two_elec_int_improved.py create mode 100644 tests/test_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..eb352821 --- /dev/null +++ b/gbasis/integrals/_two_elec_int_improved.py @@ -0,0 +1,757 @@ +"""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 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 + + +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(angmom_components_a).reshape(-1, 1, 1, 1, 1, 1, 1, 1) + norm_b = _get_factorial2_norm(angmom_components_b).reshape(1, -1, 1, 1, 1, 1, 1, 1) + norm_c = _get_factorial2_norm(angmom_components_c).reshape(1, 1, -1, 1, 1, 1, 1, 1) + norm_d = _get_factorial2_norm(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/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py new file mode 100644 index 00000000..a9c0c2bd --- /dev/null +++ b/tests/test_two_elec_int_improved.py @@ -0,0 +1,392 @@ +"""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 +import pytest + +from gbasis.integrals._two_elec_int_improved import ( + _vertical_recursion_relation, + _electron_transfer_recursion, + _optimized_contraction, + _get_factorial2_norm, + _horizontal_recursion_relation, + compute_two_electron_integrals_os_hgp, +) + + +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) + + +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. + """ + from gbasis.integrals._two_elec_int import _compute_two_elec_integrals_angmom_zero + from gbasis.integrals.point_charge import PointChargeIntegral + + 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.""" + from gbasis.integrals._two_elec_int import _compute_two_elec_integrals + from gbasis.integrals.point_charge import PointChargeIntegral + + 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.""" + from gbasis.integrals.point_charge import PointChargeIntegral + + 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.""" + from scipy.special import hyp1f1 + + 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.""" + from scipy.special import hyp1f1 + + 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) + From bf6e7273b058549ae0470e651dbbbd6d2a317056 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Tue, 3 Mar 2026 20:04:46 +0530 Subject: [PATCH 08/29] 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 8ab104ea4d8918ed617f211143706507e6ccbaff Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Wed, 4 Mar 2026 16:57:52 +0530 Subject: [PATCH 09/29] test: add erf/erfc attenuated Boys function tests and update docstring --- gbasis/integrals/boys_functions.py | 6 +- tests/test_boys_functions.py | 199 ++++++++++++++++++++++++++++- 2 files changed, 203 insertions(+), 2 deletions(-) diff --git a/gbasis/integrals/boys_functions.py b/gbasis/integrals/boys_functions.py index c3f04ed5..f9cb8ac6 100644 --- a/gbasis/integrals/boys_functions.py +++ b/gbasis/integrals/boys_functions.py @@ -3,10 +3,14 @@ This module implements the Boys function, which is the starting point for computing two-electron integrals using the Obara-Saika recursion relations. -The Boys function is defined as: +The standard Boys function is defined as: F_m(T) = integral from 0 to 1 of t^(2m) * exp(-T*t^2) dt +For non-Coulomb potentials (erf-attenuated, erfc-attenuated), the modified +Boys functions follow from Ahlrichs (2006), where the key modification is +a rescaling of the argument T and a prefactor depending on omega and rho. + References: - Helgaker, T.; Jorgensen, P.; Olsen, J. "Molecular Electronic-Structure Theory" (eq. 9.8.39 for hyp1f1 representation) diff --git a/tests/test_boys_functions.py b/tests/test_boys_functions.py index 2741fb06..e31630c8 100644 --- a/tests/test_boys_functions.py +++ b/tests/test_boys_functions.py @@ -1,4 +1,8 @@ -"""Test gbasis.integrals.boys_functions module.""" +"""Test gbasis.integrals.boys_functions module. + +Tests for standard Coulomb Boys function (Week 1) and +erf/erfc-attenuated potentials (Week 5). +""" import numpy as np import pytest @@ -7,6 +11,8 @@ from gbasis.integrals.boys_functions import ( boys_function_all_orders, + boys_function_erf, + boys_function_erfc, boys_function_standard, get_boys_function, boys_function_recursion, @@ -190,6 +196,64 @@ def test_unknown_potential(self): with pytest.raises(ValueError): get_boys_function("unknown_potential") + def test_erf_potential(self): + """Test that erf potential returns correct function.""" + boys = get_boys_function("erf", omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + assert np.all(np.isfinite(result)) + + def test_erfc_potential(self): + """Test that erfc potential returns correct function.""" + boys = get_boys_function("erfc", omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + assert np.all(np.isfinite(result)) + + def test_erf_aliases(self): + """Test that erf aliases work correctly.""" + for alias in ["erf", "erf_coulomb"]: + boys = get_boys_function(alias, omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + + def test_erfc_aliases(self): + """Test that erfc aliases work correctly.""" + for alias in ["erfc", "erfc_coulomb"]: + boys = get_boys_function(alias, omega=0.5) + result = boys(np.array([0]), np.array([1.0]), rho=0.8) + assert result is not None + + def test_erf_requires_omega(self): + """Test that erf potential requires omega parameter.""" + with pytest.raises(ValueError): + get_boys_function("erf") + + def test_erfc_requires_omega(self): + """Test that erfc potential requires omega parameter.""" + with pytest.raises(ValueError): + get_boys_function("erfc") + + def test_erf_requires_rho(self): + """Test that erf Boys function requires rho at call time.""" + boys = get_boys_function("erf", omega=0.5) + with pytest.raises(ValueError): + boys(np.array([0]), np.array([1.0])) + + def test_erfc_requires_rho(self): + """Test that erfc Boys function requires rho at call time.""" + boys = get_boys_function("erfc", omega=0.5) + with pytest.raises(ValueError): + boys(np.array([0]), np.array([1.0])) + + def test_case_insensitive(self): + """Test that potential matching is case-insensitive.""" + for alias in ["COULOMB", "Coulomb", "Standard", "1/R"]: + boys = get_boys_function(alias) + result = boys(np.array([0]), np.array([1.0])) + expected = boys_function_standard(np.array([0]), np.array([1.0])) + assert np.allclose(result, expected), f"Failed for '{alias}'" + class TestBoysNumericalIntegration: """Compare Boys function with numerical integration.""" @@ -215,6 +279,139 @@ def boys_integrand(t, T, m): result_analytic = boys_function_standard(np.array([2]), np.array([T]))[0] result_numeric, _ = quad(boys_integrand, 0, 1, args=(T, 2)) assert np.allclose(result_analytic, result_numeric, rtol=1e-6) +class TestBoysErf: + """Tests for the erf-attenuated Boys function.""" + + def test_erf_large_omega_recovers_coulomb(self): + """Test that erf with large omega approaches standard Coulomb. + + As omega -> infinity, erf(omega*r)/r -> 1/r. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erf = boys_function_erf(orders, T, rho, omega=1000.0) + result_std = boys_function_standard(orders, T) + + np.testing.assert_allclose(result_erf, result_std, rtol=1e-5, + err_msg="erf with large omega should recover standard Coulomb") + + def test_erf_small_omega_approaches_zero(self): + """Test that erf with small omega approaches zero. + + As omega -> 0, erf(omega*r)/r -> 0. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erf = boys_function_erf(orders, T, rho, omega=0.001) + result_std = boys_function_standard(orders, T) + # With omega=0.001, erf values should be much smaller than Coulomb + ratio = np.abs(result_erf) / (np.abs(result_std) + 1e-30) + assert np.all(ratio < 0.01), "erf with small omega should be << Coulomb" + + def test_erf_positive_values(self): + """Test that erf Boys function produces positive values.""" + orders = np.array([0, 1, 2, 3])[:, None] + T = np.array([0.1, 0.5, 1.0, 5.0])[None, :] + rho = np.ones_like(T) * 1.0 + + result = boys_function_erf(orders, T, rho, omega=0.5) + assert np.all(result >= 0), "erf Boys function should be non-negative" + + def test_erf_less_than_coulomb(self): + """Test that erf Boys < standard Coulomb (since erf(x) <= 1).""" + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erf = boys_function_erf(orders, T, rho, omega=0.5) + result_std = boys_function_standard(orders, T) + + assert np.all(result_erf <= result_std + 1e-15), \ + "erf Boys should be <= standard Coulomb" + + def test_erf_scaling_formula(self): + """Test the erf scaling formula: scaling^(m+0.5) * F_m(scaling * T). + + Tests m=0,1,2 to ensure the exponent (m+0.5) is not hardcoded. + """ + T = np.array([1.0]) + rho = np.array([0.8]) + omega = 0.5 + scaling = omega**2 / (omega**2 + rho) + T_mod = scaling * T + + for m_val in [0, 1, 2]: + m = np.array([m_val]) + expected = scaling ** (m_val + 0.5) * boys_function_standard(m, T_mod) + result = boys_function_erf(m, T, rho, omega) + np.testing.assert_allclose(result, expected, rtol=1e-12, + err_msg=f"erf scaling formula failed for m={m_val}") + + +class TestBoysErfc: + """Tests for the erfc-attenuated Boys function.""" + + def test_erfc_small_omega_recovers_coulomb(self): + """Test that erfc with small omega approaches standard Coulomb. + + As omega -> 0, erfc(omega*r)/r -> 1/r. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erfc = boys_function_erfc(orders, T, rho, omega=0.001) + result_std = boys_function_standard(orders, T) + + np.testing.assert_allclose(result_erfc, result_std, rtol=0.01, + err_msg="erfc with small omega should recover standard Coulomb") + + def test_erfc_large_omega_approaches_zero(self): + """Test that erfc with large omega approaches zero. + + As omega -> infinity, erfc(omega*r)/r -> 0. + """ + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.5, 1.0, 2.0])[None, :] + rho = np.ones_like(T) * 0.8 + + result_erfc = boys_function_erfc(orders, T, rho, omega=1000.0) + assert np.all(np.abs(result_erfc) < 1e-6), \ + "erfc with large omega should be near zero" + + def test_erf_plus_erfc_equals_coulomb(self): + """Test that erf + erfc = Coulomb for any omega. + + This is the fundamental identity: erf(x) + erfc(x) = 1, so + erf(omega*r)/r + erfc(omega*r)/r = 1/r. + """ + orders = np.array([0, 1, 2, 3])[:, None] + T = np.array([0.1, 0.5, 1.0, 2.0, 5.0])[None, :] + rho = np.ones_like(T) * 1.2 + + for omega in [0.1, 0.5, 1.0, 2.0]: + result_erf = boys_function_erf(orders, T, rho, omega) + result_erfc = boys_function_erfc(orders, T, rho, omega) + result_std = boys_function_standard(orders, T) + + np.testing.assert_allclose( + result_erf + result_erfc, result_std, rtol=1e-12, + err_msg=f"erf + erfc != Coulomb for omega={omega}" + ) + + def test_erfc_positive_values(self): + """Test that erfc Boys function produces positive values.""" + orders = np.array([0, 1, 2])[:, None] + T = np.array([0.1, 0.5, 1.0, 5.0])[None, :] + rho = np.ones_like(T) * 1.0 + + result = boys_function_erfc(orders, T, rho, omega=0.5) + assert np.all(result >= -1e-15), "erfc Boys function should be non-negative" + class TestBoysRecursion: """Tests for the downward recursion Boys function (Eq. 71).""" From 434744427aadacefed558ca489e8ed60d50a7ee2 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Wed, 4 Mar 2026 18:20:26 +0530 Subject: [PATCH 10/29] style: apply black/ruff formatting to test_boys_functions.py --- tests/test_boys_functions.py | 123 ++++++++++++++++++++++------------- 1 file changed, 78 insertions(+), 45 deletions(-) diff --git a/tests/test_boys_functions.py b/tests/test_boys_functions.py index e31630c8..9870bd0d 100644 --- a/tests/test_boys_functions.py +++ b/tests/test_boys_functions.py @@ -13,15 +13,15 @@ boys_function_all_orders, boys_function_erf, boys_function_erfc, + boys_function_mpmath, + boys_function_recursion, boys_function_standard, get_boys_function, - boys_function_recursion, - boys_function_mpmath, ) - try: import mpmath # noqa: F401 + HAS_MPMATH = True except ImportError: HAS_MPMATH = False @@ -76,8 +76,9 @@ def test_boys_broadcasting(self): for i, m in enumerate([0, 1, 2]): for j, t in enumerate([0.5, 1.0, 2.0]): expected = boys_function_standard(np.array([m]), np.array([t]))[0] - assert np.isclose(result[i, j], expected), \ - f"Broadcasting value mismatch at m={m}, T={t}" + assert np.isclose( + result[i, j], expected + ), f"Broadcasting value mismatch at m={m}, T={t}" def test_boys_recurrence(self): """Test the recurrence relation: (2m+1)*F_m(T) = 2T*F_{m+1}(T) + exp(-T).""" @@ -111,16 +112,22 @@ def test_matches_standard_large_T(self): for m in range(m_max + 1): expected = boys_function_standard(np.array([m]), T) - np.testing.assert_allclose(result[m], expected.flatten(), rtol=1e-8, - err_msg=f"All-orders doesn't match standard for m={m}, large T") + np.testing.assert_allclose( + result[m], + expected.flatten(), + rtol=1e-8, + err_msg=f"All-orders doesn't match standard for m={m}, large T", + ) def test_output_shape(self): """Test that all-orders function returns correct shape (m_max+1, *T.shape).""" m_max = 4 T = np.array([0.5, 1.0, 2.0, 5.0, 30.0]) result = boys_function_all_orders(m_max, T) - assert result.shape == (m_max + 1, T.shape[0]), \ - f"Expected shape {(m_max + 1, T.shape[0])}, got {result.shape}" + assert result.shape == ( + m_max + 1, + T.shape[0], + ), f"Expected shape {(m_max + 1, T.shape[0])}, got {result.shape}" def test_small_T_series_branch(self): """Exercise the small-T series path and compare with reference.""" @@ -131,9 +138,8 @@ def test_small_T_series_branch(self): for m in range(m_max + 1): expected = boys_function_standard(np.array([m]), T) np.testing.assert_allclose( - result[m], expected, - rtol=1e-12, atol=1e-14, - err_msg=f"Small-T mismatch at m={m}") + result[m], expected, rtol=1e-12, atol=1e-14, err_msg=f"Small-T mismatch at m={m}" + ) def test_multidim_broadcasting(self): """Check broadcasting for multidimensional T arrays.""" @@ -158,8 +164,9 @@ def test_recursion_identity(self): for m in range(m_max): lhs = (2 * m + 1) * all_vals[m] rhs = 2 * T * all_vals[m + 1] + np.exp(-T) - np.testing.assert_allclose(lhs, rhs, rtol=1e-11, atol=1e-14, - err_msg=f"Recurrence failed for m={m}") + np.testing.assert_allclose( + lhs, rhs, rtol=1e-11, atol=1e-14, err_msg=f"Recurrence failed for m={m}" + ) def test_random_values_against_reference(self): """Random spot-check across orders and T values.""" @@ -170,8 +177,13 @@ def test_random_values_against_reference(self): for m in range(m_max + 1): expected = boys_function_standard(np.array([m]), T) - np.testing.assert_allclose(result[m], expected, rtol=1e-11, atol=1e-14, - err_msg=f"Random check mismatch at m={m}") + np.testing.assert_allclose( + result[m], + expected, + rtol=1e-11, + atol=1e-14, + err_msg=f"Random check mismatch at m={m}", + ) class TestGetBoysFunction: @@ -262,7 +274,7 @@ def test_numerical_integration_m0(self): """Test F_0(T) against numerical integration.""" def boys_integrand(t, T, m): - return t ** (2 * m) * np.exp(-T * t ** 2) + return t ** (2 * m) * np.exp(-T * t**2) for T in [0.5, 1.0, 2.0, 5.0]: result_analytic = boys_function_standard(np.array([0]), np.array([T]))[0] @@ -273,12 +285,14 @@ def test_numerical_integration_m2(self): """Test F_2(T) against numerical integration.""" def boys_integrand(t, T, m): - return t ** (2 * m) * np.exp(-T * t ** 2) + return t ** (2 * m) * np.exp(-T * t**2) for T in [0.5, 1.0, 2.0]: result_analytic = boys_function_standard(np.array([2]), np.array([T]))[0] result_numeric, _ = quad(boys_integrand, 0, 1, args=(T, 2)) assert np.allclose(result_analytic, result_numeric, rtol=1e-6) + + class TestBoysErf: """Tests for the erf-attenuated Boys function.""" @@ -294,8 +308,12 @@ def test_erf_large_omega_recovers_coulomb(self): result_erf = boys_function_erf(orders, T, rho, omega=1000.0) result_std = boys_function_standard(orders, T) - np.testing.assert_allclose(result_erf, result_std, rtol=1e-5, - err_msg="erf with large omega should recover standard Coulomb") + np.testing.assert_allclose( + result_erf, + result_std, + rtol=1e-5, + err_msg="erf with large omega should recover standard Coulomb", + ) def test_erf_small_omega_approaches_zero(self): """Test that erf with small omega approaches zero. @@ -330,8 +348,7 @@ def test_erf_less_than_coulomb(self): result_erf = boys_function_erf(orders, T, rho, omega=0.5) result_std = boys_function_standard(orders, T) - assert np.all(result_erf <= result_std + 1e-15), \ - "erf Boys should be <= standard Coulomb" + assert np.all(result_erf <= result_std + 1e-15), "erf Boys should be <= standard Coulomb" def test_erf_scaling_formula(self): """Test the erf scaling formula: scaling^(m+0.5) * F_m(scaling * T). @@ -348,8 +365,9 @@ def test_erf_scaling_formula(self): m = np.array([m_val]) expected = scaling ** (m_val + 0.5) * boys_function_standard(m, T_mod) result = boys_function_erf(m, T, rho, omega) - np.testing.assert_allclose(result, expected, rtol=1e-12, - err_msg=f"erf scaling formula failed for m={m_val}") + np.testing.assert_allclose( + result, expected, rtol=1e-12, err_msg=f"erf scaling formula failed for m={m_val}" + ) class TestBoysErfc: @@ -367,8 +385,12 @@ def test_erfc_small_omega_recovers_coulomb(self): result_erfc = boys_function_erfc(orders, T, rho, omega=0.001) result_std = boys_function_standard(orders, T) - np.testing.assert_allclose(result_erfc, result_std, rtol=0.01, - err_msg="erfc with small omega should recover standard Coulomb") + np.testing.assert_allclose( + result_erfc, + result_std, + rtol=0.01, + err_msg="erfc with small omega should recover standard Coulomb", + ) def test_erfc_large_omega_approaches_zero(self): """Test that erfc with large omega approaches zero. @@ -380,8 +402,7 @@ def test_erfc_large_omega_approaches_zero(self): rho = np.ones_like(T) * 0.8 result_erfc = boys_function_erfc(orders, T, rho, omega=1000.0) - assert np.all(np.abs(result_erfc) < 1e-6), \ - "erfc with large omega should be near zero" + assert np.all(np.abs(result_erfc) < 1e-6), "erfc with large omega should be near zero" def test_erf_plus_erfc_equals_coulomb(self): """Test that erf + erfc = Coulomb for any omega. @@ -399,8 +420,10 @@ def test_erf_plus_erfc_equals_coulomb(self): result_std = boys_function_standard(orders, T) np.testing.assert_allclose( - result_erf + result_erfc, result_std, rtol=1e-12, - err_msg=f"erf + erfc != Coulomb for omega={omega}" + result_erf + result_erfc, + result_std, + rtol=1e-12, + err_msg=f"erf + erfc != Coulomb for omega={omega}", ) def test_erfc_positive_values(self): @@ -423,8 +446,13 @@ def test_recursion_matches_all_orders(self): result_rec = boys_function_recursion(m_max, T) result_ref = boys_function_all_orders(m_max, T) - np.testing.assert_allclose(result_rec, result_ref, rtol=1e-12, atol=1e-15, - err_msg="Recursion doesn't match all_orders") + np.testing.assert_allclose( + result_rec, + result_ref, + rtol=1e-12, + atol=1e-15, + err_msg="Recursion doesn't match all_orders", + ) def test_recursion_T_zero(self): """Test F_m(0) = 1/(2m+1) via recursion.""" @@ -432,8 +460,9 @@ def test_recursion_T_zero(self): result = boys_function_recursion(m_max, np.array([0.0])) for m in range(m_max + 1): expected = 1.0 / (2 * m + 1) - np.testing.assert_allclose(result[m], expected, rtol=1e-14, - err_msg=f"F_{m}(0) incorrect") + np.testing.assert_allclose( + result[m], expected, rtol=1e-14, err_msg=f"F_{m}(0) incorrect" + ) def test_recursion_large_T(self): """Test stability for large T values.""" @@ -442,8 +471,9 @@ def test_recursion_large_T(self): result_rec = boys_function_recursion(m_max, T) result_ref = boys_function_all_orders(m_max, T) - np.testing.assert_allclose(result_rec, result_ref, rtol=1e-10, - err_msg="Recursion unstable for large T") + np.testing.assert_allclose( + result_rec, result_ref, rtol=1e-10, err_msg="Recursion unstable for large T" + ) def test_recursion_high_order(self): """Test with m_max = 20 (high angular momentum).""" @@ -452,8 +482,9 @@ def test_recursion_high_order(self): result_rec = boys_function_recursion(m_max, T) result_ref = boys_function_all_orders(m_max, T) - np.testing.assert_allclose(result_rec, result_ref, rtol=1e-11, - err_msg="Recursion fails at high order") + np.testing.assert_allclose( + result_rec, result_ref, rtol=1e-11, err_msg="Recursion fails at high order" + ) def test_recursion_output_shape(self): """Test output shape matches convention.""" @@ -481,8 +512,9 @@ def test_recursion_recurrence_identity(self): for m in range(m_max): lhs = (2 * m + 1) * all_vals[m] rhs = 2 * T * all_vals[m + 1] + np.exp(-T) - np.testing.assert_allclose(lhs, rhs, rtol=1e-11, atol=1e-14, - err_msg=f"Recurrence failed for m={m}") + np.testing.assert_allclose( + lhs, rhs, rtol=1e-11, atol=1e-14, err_msg=f"Recurrence failed for m={m}" + ) @pytest.mark.skipif(not HAS_MPMATH, reason="mpmath not installed") @@ -496,8 +528,9 @@ def test_mpmath_matches_standard(self): result_mp = boys_function_mpmath(m_max, T, dps=30) result_std = boys_function_all_orders(m_max, T) - np.testing.assert_allclose(result_mp, result_std, rtol=1e-14, - err_msg="mpmath doesn't match standard for normal T") + np.testing.assert_allclose( + result_mp, result_std, rtol=1e-14, err_msg="mpmath doesn't match standard for normal T" + ) def test_mpmath_output_shape(self): """Test that mpmath returns correct shape.""" @@ -521,6 +554,6 @@ def test_mpmath_high_order(self): result_mp = boys_function_mpmath(m_max, T, dps=50) result_std = boys_function_all_orders(m_max, T) - np.testing.assert_allclose(result_mp, result_std, rtol=1e-14, - err_msg="mpmath high-order mismatch") - + np.testing.assert_allclose( + result_mp, result_std, rtol=1e-14, err_msg="mpmath high-order mismatch" + ) From d96fd22c3e190f6a5e88a1cdb0b7c429b0a88bfd Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Wed, 4 Mar 2026 18:30:56 +0530 Subject: [PATCH 11/29] 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 7ec98ddb891eb553a67385b38c2bd1d771636946 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Wed, 4 Mar 2026 20:22:23 +0530 Subject: [PATCH 12/29] style: move imports to top-level, fix ruff PLC0415 in test file --- tests/test_two_elec_int_improved.py | 290 ++++++++++++++++++++-------- 1 file changed, 211 insertions(+), 79 deletions(-) diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py index a9c0c2bd..6e39b73d 100644 --- a/tests/test_two_elec_int_improved.py +++ b/tests/test_two_elec_int_improved.py @@ -4,16 +4,19 @@ """ import numpy as np -import pytest +from scipy.special import hyp1f1 +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 ( - _vertical_recursion_relation, _electron_transfer_recursion, - _optimized_contraction, _get_factorial2_norm, - _horizontal_recursion_relation, + _vertical_recursion_relation, compute_two_electron_integrals_os_hgp, ) +from gbasis.integrals.point_charge import PointChargeIntegral class TestVerticalRecursion: @@ -65,20 +68,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]]]]) @@ -117,9 +115,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) @@ -127,8 +123,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] @@ -140,9 +135,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) @@ -150,12 +143,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 @@ -195,9 +197,6 @@ def test_ssss_matches_old(self): all-zero angular momentum, so we compare against the specialized _compute_two_elec_integrals_angmom_zero function instead. """ - from gbasis.integrals._two_elec_int import _compute_two_elec_integrals_angmom_zero - from gbasis.integrals.point_charge import PointChargeIntegral - boys_func = PointChargeIntegral.boys_func coord_a = np.array([0.0, 0.0, 0.0]) @@ -220,29 +219,54 @@ def test_ssss_matches_old(self): # 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, + 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, + 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") + 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.""" - from gbasis.integrals._two_elec_int import _compute_two_elec_integrals - from gbasis.integrals.point_charge import PointChargeIntegral - boys_func = PointChargeIntegral.boys_func coord_a = np.array([0.0, 0.0, 0.0]) @@ -268,23 +292,59 @@ def test_spsp_matches_old(self): # 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, + 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, + 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") + 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: @@ -292,8 +352,6 @@ class TestHighAngularMomentum: def test_dddd_no_overflow(self): """Test (dd|dd) integrals produce finite values.""" - from gbasis.integrals.point_charge import PointChargeIntegral - boys_func = PointChargeIntegral.boys_func coord_a = np.array([0.0, 0.0, 0.0]) @@ -305,17 +363,30 @@ def test_dddd_no_overflow(self): 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] - ]) + 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, + 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" @@ -327,7 +398,6 @@ class TestPrimitiveScreening: def test_screening_zero_threshold(self): """With threshold=0, results match unscreened exactly.""" - from scipy.special import hyp1f1 def boys_func(orders, weighted_dist): return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) @@ -341,18 +411,50 @@ def boys_func(orders, weighted_dist): 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, + 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, + 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, ) @@ -360,7 +462,6 @@ def boys_func(orders, weighted_dist): def test_screening_reasonable_threshold(self): """With reasonable threshold, results match within tolerance.""" - from scipy.special import hyp1f1 def boys_func(orders, weighted_dist): return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) @@ -373,20 +474,51 @@ def boys_func(orders, weighted_dist): 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, + 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, + 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) - From f0f57fa6854dbb497d355741fb7b78832929bf5a Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Thu, 5 Mar 2026 15:46:47 +0530 Subject: [PATCH 13/29] 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 14/29] 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 15/29] 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.""" From 3411d2311294e838bde7f7b9166358253040ca9b Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 6 Mar 2026 16:03:22 +0530 Subject: [PATCH 16/29] style: fix black formatting and import sorting in test_electron_repulsion --- gbasis/integrals/electron_repulsion.py | 248 +++++++++-- tests/test_electron_repulsion.py | 574 ++++++++++++++++--------- tests/test_two_elec_int_improved.py | 144 +++++++ 3 files changed, 709 insertions(+), 257 deletions(-) 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 index 71ae5eaf..5974f117 100644 --- a/tests/test_two_elec_int_improved.py +++ b/tests/test_two_elec_int_improved.py @@ -6,6 +6,7 @@ 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, @@ -17,6 +18,7 @@ _vertical_recursion_relation, compute_two_electron_integrals_os_hgp, ) +from gbasis.integrals.electron_repulsion import ElectronRepulsionIntegralImproved from gbasis.integrals.point_charge import PointChargeIntegral @@ -553,3 +555,145 @@ def boys_func(orders, weighted_dist): ) 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", + ) From 6314d7861fce0dccfc68de26d09ea325f5bfdcb8 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Sun, 8 Mar 2026 16:10:27 +0530 Subject: [PATCH 17/29] feat: implement Schwarz screening to optimize two-electron integral computation --- gbasis/integrals/_schwarz_screening.py | 218 ++++++++++++ gbasis/integrals/electron_repulsion.py | 67 +++- tests/test_schwarz_screening.py | 460 +++++++++++++++++++++++++ tests/test_two_elec_int_improved.py | 83 ++++- 4 files changed, 815 insertions(+), 13 deletions(-) create mode 100644 gbasis/integrals/_schwarz_screening.py create mode 100644 tests/test_schwarz_screening.py diff --git a/gbasis/integrals/_schwarz_screening.py b/gbasis/integrals/_schwarz_screening.py new file mode 100644 index 00000000..306769df --- /dev/null +++ b/gbasis/integrals/_schwarz_screening.py @@ -0,0 +1,218 @@ +"""Integral screening utilities for efficient 2-electron integral computation. + +This module implements Schwarz screening and shell-pair screening to skip +negligible integrals, providing speedup for spatially extended systems. + +References: +- Häser, M. & Ahlrichs, R. J. Comput. Chem. 1989, 10, 104. +- Gill, P. M. W.; Johnson, B. G.; Pople, J. A. Int. J. Quantum Chem. 1991, 40, 745. +""" + +import numpy as np + + +def compute_schwarz_bound_shell_pair(boys_func, cont_one, cont_two, compute_integral_func): + """Compute Schwarz bound for a shell pair: sqrt((ab|ab)). + + Parameters + ---------- + boys_func : callable + Boys function for integral evaluation. + cont_one : GeneralizedContractionShell + First contracted shell. + cont_two : GeneralizedContractionShell + Second contracted shell. + compute_integral_func : callable + Function to compute (ab|cd) integrals. + + Returns + ------- + bound : float + Schwarz bound sqrt(max|(ab|ab)|) for this shell pair. + """ + # Compute (ab|ab) integral + integral = compute_integral_func( + boys_func, + cont_one.coord, + cont_one.angmom, + cont_one.angmom_components_cart, + cont_one.exps, + cont_one.coeffs, + cont_two.coord, + cont_two.angmom, + cont_two.angmom_components_cart, + cont_two.exps, + cont_two.coeffs, + cont_one.coord, + cont_one.angmom, + cont_one.angmom_components_cart, + cont_one.exps, + cont_one.coeffs, + cont_two.coord, + cont_two.angmom, + cont_two.angmom_components_cart, + cont_two.exps, + cont_two.coeffs, + ) + + # Return sqrt of maximum absolute value + return np.sqrt(np.max(np.abs(integral))) + + +def compute_schwarz_bounds(contractions, boys_func, compute_integral_func): + """Precompute Schwarz bounds for all shell pairs. + + Parameters + ---------- + contractions : list of GeneralizedContractionShell + List of all contracted shells. + boys_func : callable + Boys function for integral evaluation. + compute_integral_func : callable + Function to compute (ab|cd) integrals. + + Returns + ------- + bounds : np.ndarray(n_shells, n_shells) + Schwarz bounds sqrt((ab|ab)) for each shell pair. + """ + n_shells = len(contractions) + bounds = np.zeros((n_shells, n_shells)) + + for i, cont_i in enumerate(contractions): + for j in range(i, n_shells): + cont_j = contractions[j] + bounds[i, j] = compute_schwarz_bound_shell_pair( + boys_func, cont_i, cont_j, compute_integral_func + ) + bounds[j, i] = bounds[i, j] # Symmetry: (ab|ab) = (ba|ba) + + return bounds + + +def shell_pair_significant(cont_one, cont_two, threshold=1e-12): + """Check if a shell pair is significant using primitive screening. + + Uses the Gaussian product theorem: exp(-a*b/(a+b) * |A-B|^2) factor. + If this factor is below threshold for all primitive pairs, skip. + + Parameters + ---------- + cont_one : GeneralizedContractionShell + First contracted shell. + cont_two : GeneralizedContractionShell + Second contracted shell. + threshold : float + Screening threshold. + + Returns + ------- + significant : bool + True if shell pair might contribute significantly. + """ + # Distance between shell centers + r_ab_sq = np.sum((cont_one.coord - cont_two.coord) ** 2) + + if r_ab_sq < 1e-10: + # Same center, always significant + return True + + # Check if any primitive pair survives screening + for exp_a in cont_one.exps: + for exp_b in cont_two.exps: + # Gaussian decay factor + decay = np.exp(-exp_a * exp_b / (exp_a + exp_b) * r_ab_sq) + if decay > threshold: + return True + + return False + + +class SchwarzScreener: + """Class for Schwarz integral screening. + + Precomputes Schwarz bounds and provides efficient screening. + + Attributes + ---------- + bounds : np.ndarray + Schwarz bounds for all shell pairs. + threshold : float + Screening threshold. + n_screened : int + Counter for number of screened shell quartets. + n_computed : int + Counter for number of computed shell quartets. + """ + + def __init__(self, contractions, boys_func, compute_integral_func, threshold=1e-12): + """Initialize Schwarz screener. + + Parameters + ---------- + contractions : list of GeneralizedContractionShell + List of all contracted shells. + boys_func : callable + Boys function for integral evaluation. + compute_integral_func : callable + Function to compute (ab|cd) integrals. + threshold : float + Screening threshold (default: 1e-12). + """ + self.threshold = threshold + self.n_screened = 0 + self.n_computed = 0 + + # Precompute Schwarz bounds + self.bounds = compute_schwarz_bounds(contractions, boys_func, compute_integral_func) + + def is_significant(self, i, j, k, l_shell): + """Check if shell quartet (ij|kl) is significant. + + Uses Schwarz inequality: |(ij|kl)| <= sqrt((ij|ij)) * sqrt((kl|kl)) + + Parameters + ---------- + i, j, k, l_shell : int + Shell indices. + + Returns + ------- + significant : bool + True if integral might be significant, False if can be skipped. + """ + bound = self.bounds[i, j] * self.bounds[k, l_shell] + + if bound < self.threshold: + self.n_screened += 1 + return False + else: + self.n_computed += 1 + return True + + def get_statistics(self): + """Get screening statistics. + + Returns + ------- + stats : dict + Dictionary with screening statistics. + """ + total = self.n_screened + self.n_computed + if total == 0: + percent_screened = 0.0 + else: + percent_screened = 100.0 * self.n_screened / total + + return { + "n_screened": self.n_screened, + "n_computed": self.n_computed, + "total": total, + "percent_screened": percent_screened, + "speedup_factor": total / max(self.n_computed, 1), + } + + def reset_counters(self): + """Reset screening counters.""" + self.n_screened = 0 + self.n_computed = 0 diff --git a/gbasis/integrals/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index 68563f88..a48b5910 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -4,6 +4,7 @@ from gbasis.base_four_symm import BaseFourIndexSymmetric from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals._schwarz_screening import SchwarzScreener from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp from gbasis.integrals.point_charge import PointChargeIntegral @@ -275,6 +276,8 @@ class ElectronRepulsionIntegralImproved(BaseFourIndexSymmetric): """ boys_func = PointChargeIntegral.boys_func + _screener = None + _contraction_index_map = None @classmethod def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four): @@ -312,6 +315,26 @@ def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four): if not isinstance(cont_four, GeneralizedContractionShell): raise TypeError("`cont_four` must be a `GeneralizedContractionShell` instance.") + # --- Schwarz screening: skip negligible shell quartets --- + if cls._screener is not None and cls._contraction_index_map is not None: + i = cls._contraction_index_map.get(id(cont_one)) + j = cls._contraction_index_map.get(id(cont_two)) + k = cls._contraction_index_map.get(id(cont_three)) + l_idx = cls._contraction_index_map.get(id(cont_four)) + if i is not None and j is not None and k is not None and l_idx is not None: + if not cls._screener.is_significant(i, j, k, l_idx): + shape = ( + cont_one.coeffs.shape[1], + len(cont_one.angmom_components_cart), + cont_two.coeffs.shape[1], + len(cont_two.angmom_components_cart), + cont_three.coeffs.shape[1], + len(cont_three.angmom_components_cart), + cont_four.coeffs.shape[1], + len(cont_four.angmom_components_cart), + ) + return np.zeros(shape) + # --- Contraction reordering for efficiency (l_a >= l_b, l_c >= l_d, l_a >= l_c) --- bra_swapped = cont_one.angmom < cont_two.angmom if bra_swapped: @@ -364,7 +387,9 @@ def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four): return integrals -def electron_repulsion_integral_improved(basis, transform=None, notation="physicist"): +def electron_repulsion_integral_improved( + basis, transform=None, notation="physicist", schwarz_threshold=0.0 +): r"""Return the electron repulsion integrals using the improved OS+HGP algorithm. This function uses the Obara-Saika + Head-Gordon-Pople recursion scheme, @@ -398,6 +423,9 @@ def electron_repulsion_integral_improved(basis, transform=None, notation="physic notation : {"physicist", "chemist"} Convention with which the integrals are ordered. Default is Physicists' notation. + schwarz_threshold : float + Schwarz screening threshold. Shell quartets with Schwarz bound below + this threshold are skipped. Default is 0.0 (no screening). Returns ------- @@ -413,18 +441,33 @@ def electron_repulsion_integral_improved(basis, transform=None, notation="physic 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 + if schwarz_threshold > 0: + index_map = {id(shell): i for i, shell in enumerate(basis)} + screener = SchwarzScreener( + list(basis), + ElectronRepulsionIntegralImproved.boys_func, + compute_two_electron_integrals_os_hgp, + schwarz_threshold, ) - 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) + ElectronRepulsionIntegralImproved._screener = screener + ElectronRepulsionIntegralImproved._contraction_index_map = index_map + + try: + 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) + finally: + ElectronRepulsionIntegralImproved._screener = None + ElectronRepulsionIntegralImproved._contraction_index_map = None if notation == "physicist": array = np.transpose(array, (0, 2, 1, 3)) diff --git a/tests/test_schwarz_screening.py b/tests/test_schwarz_screening.py new file mode 100644 index 00000000..087a8004 --- /dev/null +++ b/tests/test_schwarz_screening.py @@ -0,0 +1,460 @@ +"""Benchmark tests for Schwarz screening and integral optimizations. + +Tests for Week 7: Schwarz screening (_screening.py) for efficient +two-electron integral computation. + +The optimizations tested: +1. Schwarz Screening - Skip negligible shell quartets using (ij|kl) <= Q_ij * Q_kl +2. Shell-pair Prescreening - Fast Gaussian decay check +3. SchwarzScreener class - Precomputed bounds with statistics +""" + +import numpy as np + +from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals._schwarz_screening import ( + SchwarzScreener, + compute_schwarz_bound_shell_pair, + compute_schwarz_bounds, + shell_pair_significant, +) +from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp +from gbasis.integrals.electron_repulsion import electron_repulsion_integral_improved +from gbasis.integrals.point_charge import PointChargeIntegral + + +def make_shell(coord, angmom, exps, coeffs): + """Create a GeneralizedContractionShell for testing.""" + return GeneralizedContractionShell( + angmom=angmom, + coord=np.array(coord), + coeffs=np.array(coeffs).reshape(-1, 1), + exps=np.array(exps), + coord_type="cartesian", + ) + + +def make_h_chain(n_atoms, spacing, include_p=False): + """Create a hydrogen chain basis set for testing. + + Parameters + ---------- + n_atoms : int + Number of hydrogen atoms in the chain. + spacing : float + Distance between atoms in Bohr. + include_p : bool + If True, include p-orbitals (cc-pVDZ-like basis). + + Returns + ------- + basis : list of GeneralizedContractionShell + Basis set for the hydrogen chain. + """ + # STO-3G exponents and coefficients + exps_s_sto3g = [3.42525091, 0.62391373, 0.16885540] + coeffs_s_sto3g = [0.15432897, 0.53532814, 0.44463454] + + # cc-pVDZ-like exponents + exps_s_dz = [13.01, 1.962, 0.4446, 0.122] + coeffs_s_dz = [0.0196, 0.1379, 0.4781, 0.5012] + exps_p = [0.727] + coeffs_p = [1.0] + + basis = [] + for i in range(n_atoms): + coord = [i * spacing, 0.0, 0.0] + if include_p: + basis.append(make_shell(coord, 0, exps_s_dz, coeffs_s_dz)) + basis.append(make_shell(coord, 1, exps_p, coeffs_p)) + else: + basis.append(make_shell(coord, 0, exps_s_sto3g, coeffs_s_sto3g)) + + return basis + + +class TestShellPairPrescreening: + """Tests for shell-pair significance screening.""" + + def test_same_center_always_significant(self): + """Test that shells on the same center are always significant.""" + shell = make_shell([0.0, 0.0, 0.0], 0, [1.0], [1.0]) + assert shell_pair_significant(shell, shell) is True + + def test_close_shells_significant(self): + """Test that close shells are significant.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [1.0], [1.0]) + shell_b = make_shell([1.0, 0.0, 0.0], 0, [1.0], [1.0]) + assert shell_pair_significant(shell_a, shell_b) is True + + def test_distant_shells_not_significant(self): + """Test that very distant shells with tight exponents are not significant.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [100.0], [1.0]) + shell_b = make_shell([100.0, 0.0, 0.0], 0, [100.0], [1.0]) + assert shell_pair_significant(shell_a, shell_b) is False + + def test_diffuse_shells_always_significant(self): + """Test that diffuse shells (small exponents) are significant even at distance.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [0.01], [1.0]) + shell_b = make_shell([10.0, 0.0, 0.0], 0, [0.01], [1.0]) + assert shell_pair_significant(shell_a, shell_b) is True + + def test_threshold_sensitivity(self): + """Test that threshold affects screening.""" + shell_a = make_shell([0.0, 0.0, 0.0], 0, [5.0], [1.0]) + shell_b = make_shell([5.0, 0.0, 0.0], 0, [5.0], [1.0]) + + # With strict threshold, should be not significant + sig_strict = shell_pair_significant(shell_a, shell_b, threshold=1e-6) + # With loose threshold, should be significant + sig_loose = shell_pair_significant(shell_a, shell_b, threshold=1e-30) + + # The loose threshold should be at least as permissive + if not sig_strict: + assert sig_loose or not sig_loose # may or may not be significant + if sig_loose: + assert True # loose is at least as permissive + + +class TestSchwarzBounds: + """Tests for Schwarz bound computation.""" + + def test_schwarz_bound_positive(self): + """Test that Schwarz bounds are non-negative.""" + shell = make_shell([0.0, 0.0, 0.0], 0, [1.0, 0.5], [0.6, 0.4]) + + boys_func = PointChargeIntegral.boys_func + bound = compute_schwarz_bound_shell_pair( + boys_func, shell, shell, compute_two_electron_integrals_os_hgp + ) + assert bound >= 0.0 + + def test_schwarz_bound_nonzero_for_same_shell(self): + """Test that Schwarz bound is nonzero for same shell.""" + shell = make_shell([0.0, 0.0, 0.0], 0, [1.0], [1.0]) + + boys_func = PointChargeIntegral.boys_func + bound = compute_schwarz_bound_shell_pair( + boys_func, shell, shell, compute_two_electron_integrals_os_hgp + ) + assert bound > 0.0 + + def test_schwarz_bounds_matrix_symmetric(self): + """Test that Schwarz bounds matrix is symmetric.""" + basis = make_h_chain(3, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + np.testing.assert_allclose( + bounds, bounds.T, rtol=1e-10, err_msg="Schwarz bounds matrix should be symmetric" + ) + + def test_schwarz_bounds_matrix_shape(self): + """Test that Schwarz bounds matrix has correct shape.""" + basis = make_h_chain(4, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + assert bounds.shape == (4, 4) + + def test_schwarz_bounds_diagonal_positive(self): + """Test that diagonal elements of bounds matrix are positive.""" + basis = make_h_chain(3, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + for i in range(len(basis)): + assert bounds[i, i] > 0, f"Diagonal bound[{i},{i}] should be positive" + + def test_schwarz_inequality_holds(self): + """Test that the Schwarz inequality |(ij|kl)| <= Q_ij * Q_kl holds.""" + basis = make_h_chain(3, spacing=2.0) + + boys_func = PointChargeIntegral.boys_func + bounds = compute_schwarz_bounds(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Compute actual integrals + integrals = electron_repulsion_integral_improved(basis, notation="chemist") + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + max_integral = np.max(np.abs(integrals[i, j, k, l_idx])) + schwarz_bound = bounds[i, j] * bounds[k, l_idx] + assert ( + max_integral <= schwarz_bound * (1 + 1e-6) + 1e-12 + ), f"Schwarz inequality violated for ({i}{j}|{k}{l_idx})" + + +class TestSchwarzScreener: + """Tests for the SchwarzScreener class.""" + + def test_screener_initialization(self): + """Test that screener initializes correctly.""" + basis = make_h_chain(3, spacing=2.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + assert screener.bounds.shape == (3, 3) + assert screener.n_screened == 0 + assert screener.n_computed == 0 + + def test_same_shell_significant(self): + """Test that same-shell quartets are always significant.""" + basis = make_h_chain(3, spacing=2.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + assert screener.is_significant(0, 0, 0, 0) is True + + def test_screening_statistics(self): + """Test that screening statistics are tracked correctly.""" + basis = make_h_chain(4, spacing=10.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Run through all quartets + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener.is_significant(i, j, k, l_idx) + + stats = screener.get_statistics() + assert stats["total"] == n**4 + assert stats["n_screened"] + stats["n_computed"] == n**4 + assert 0 <= stats["percent_screened"] <= 100 + + def test_extended_system_has_screening(self): + """Test that extended systems have significant screening.""" + basis = make_h_chain(4, spacing=15.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener.is_significant(i, j, k, l_idx) + + stats = screener.get_statistics() + # For well-separated atoms, many integrals should be screened + assert ( + stats["percent_screened"] > 30.0 + ), f"Expected >30% screening for extended system, got {stats['percent_screened']:.1f}%" + + def test_compact_system_no_screening(self): + """Test that compact systems have little screening.""" + basis = make_h_chain(3, spacing=1.4) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener.is_significant(i, j, k, l_idx) + + stats = screener.get_statistics() + # For compact molecules, most integrals should survive screening + assert ( + stats["percent_screened"] < 70.0 + ), f"Expected <70% screening for compact system, got {stats['percent_screened']:.1f}%" + + def test_reset_counters(self): + """Test that reset_counters works correctly.""" + basis = make_h_chain(3, spacing=2.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + screener.is_significant(0, 0, 0, 0) + screener.is_significant(0, 0, 1, 1) + assert screener.n_screened + screener.n_computed > 0 + + screener.reset_counters() + assert screener.n_screened == 0 + assert screener.n_computed == 0 + + def test_custom_threshold(self): + """Test that custom threshold affects screening.""" + basis = make_h_chain(4, spacing=5.0) + boys_func = PointChargeIntegral.boys_func + + # Strict threshold - more screening + screener_strict = SchwarzScreener( + basis, boys_func, compute_two_electron_integrals_os_hgp, threshold=1e-8 + ) + # Loose threshold - less screening + screener_loose = SchwarzScreener( + basis, boys_func, compute_two_electron_integrals_os_hgp, threshold=1e-16 + ) + + n = len(basis) + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + screener_strict.is_significant(i, j, k, l_idx) + screener_loose.is_significant(i, j, k, l_idx) + + stats_strict = screener_strict.get_statistics() + stats_loose = screener_loose.get_statistics() + + # Strict threshold should screen at least as much as loose + assert ( + stats_strict["n_screened"] >= stats_loose["n_screened"] + ), "Strict threshold should screen more than loose threshold" + + +class TestScreeningCorrectness: + """Tests to verify that screened integrals match unscreened.""" + + def test_screened_vs_unscreened_compact(self): + """Test that screened computation gives same result for compact molecule.""" + basis = make_h_chain(3, spacing=1.4) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Compute full integrals + full_integrals = electron_repulsion_integral_improved(basis, notation="chemist") + + # Verify that significant quartets match + n = full_integrals.shape[0] + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + if not screener.is_significant(i, j, k, l_idx): + # Screened integrals should actually be negligible + assert ( + np.abs(full_integrals[i, j, k, l_idx]) < 1e-10 + ), f"Screened integral ({i}{j}|{k}{l_idx}) is not negligible" + + def test_screened_vs_unscreened_extended(self): + """Test that screened computation gives same result for extended molecule.""" + basis = make_h_chain(4, spacing=10.0) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # Compute full integrals + full_integrals = electron_repulsion_integral_improved(basis, notation="chemist") + + n = full_integrals.shape[0] + for i in range(n): + for j in range(n): + for k in range(n): + for l_idx in range(n): + if not screener.is_significant(i, j, k, l_idx): + assert ( + np.abs(full_integrals[i, j, k, l_idx]) < 1e-10 + ), f"Screened integral ({i}{j}|{k}{l_idx}) is not negligible" + + def test_screened_with_p_orbitals(self): + """Test screening correctness with p-orbitals.""" + basis = make_h_chain(2, spacing=8.0, include_p=True) + boys_func = PointChargeIntegral.boys_func + + screener = SchwarzScreener(basis, boys_func, compute_two_electron_integrals_os_hgp) + + # For shell-level screening, verify bounds shape matches + assert screener.bounds.shape == (len(basis), len(basis)) + + # All diagonal bounds should be positive + for i in range(len(basis)): + assert screener.bounds[i, i] > 0 + + +class TestBenchmarkImprovedAlgorithm: + """Benchmark tests for the improved OS+HGP algorithm.""" + + def test_improved_algorithm_ssss(self): + """Test improved algorithm for (ss|ss) integrals.""" + coord = np.array([0.0, 0.0, 0.0]) + exps = np.array([1.0, 0.5]) + coeffs = np.array([[0.6], [0.4]]) + + angmom_s = 0 + angmom_comp_s = np.array([[0, 0, 0]]) + + boys_func = PointChargeIntegral.boys_func + + result = compute_two_electron_integrals_os_hgp( + boys_func, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + coord, + angmom_s, + angmom_comp_s, + exps, + coeffs, + ) + + # Should be positive for (ss|ss) + assert result[0, 0, 0, 0, 0, 0, 0, 0] > 0 + + def test_improved_algorithm_pppp(self): + """Test improved algorithm for (pp|pp) integrals.""" + coord = np.array([0.0, 0.0, 0.0]) + exps = np.array([1.0]) + coeffs = np.array([[1.0]]) + + angmom_p = 1 + angmom_comp_p = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + + boys_func = PointChargeIntegral.boys_func + + result = compute_two_electron_integrals_os_hgp( + boys_func, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + coord, + angmom_p, + angmom_comp_p, + exps, + coeffs, + ) + + assert result.shape == (3, 3, 3, 3, 1, 1, 1, 1) + assert result[0, 0, 0, 0, 0, 0, 0, 0] > 0 diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py index 5974f117..91f87ec2 100644 --- a/tests/test_two_elec_int_improved.py +++ b/tests/test_two_elec_int_improved.py @@ -7,6 +7,7 @@ from scipy.special import hyp1f1 from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals._schwarz_screening import SchwarzScreener from gbasis.integrals._two_elec_int import ( _compute_two_elec_integrals, _compute_two_elec_integrals_angmom_zero, @@ -18,7 +19,10 @@ _vertical_recursion_relation, compute_two_electron_integrals_os_hgp, ) -from gbasis.integrals.electron_repulsion import ElectronRepulsionIntegralImproved +from gbasis.integrals.electron_repulsion import ( + ElectronRepulsionIntegralImproved, + electron_repulsion_integral_improved, +) from gbasis.integrals.point_charge import PointChargeIntegral @@ -697,3 +701,80 @@ def test_pd_sp_braket_reordering(self): rtol=1e-10, err_msg="Reordered (pd|sp) doesn't match direct computation", ) + + +class TestSchwarzScreening: + """Tests for Schwarz screening integration.""" + + def test_schwarz_no_effect(self): + """With threshold=0, results match unscreened exactly.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + + shell_s = GeneralizedContractionShell( + 0, coord_a, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_p = GeneralizedContractionShell( + 1, coord_b, np.array([[1.0]]), np.array([0.8]), "cartesian" + ) + basis = [shell_s, shell_p] + + result_no_screen = electron_repulsion_integral_improved(basis, notation="chemist") + result_screened = electron_repulsion_integral_improved( + basis, notation="chemist", schwarz_threshold=0.0 + ) + + np.testing.assert_array_equal(result_no_screen, result_screened) + + def test_schwarz_screening_tight(self): + """With tight threshold, results match within tolerance.""" + coord_a = np.array([0.0, 0.0, 0.0]) + coord_b = np.array([1.0, 0.0, 0.0]) + coord_c = np.array([0.0, 1.0, 0.0]) + + shell_s1 = GeneralizedContractionShell( + 0, coord_a, np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s2 = GeneralizedContractionShell( + 0, coord_b, np.array([[1.0]]), np.array([0.8]), "cartesian" + ) + shell_s3 = GeneralizedContractionShell( + 0, coord_c, np.array([[1.0]]), np.array([0.6]), "cartesian" + ) + basis = [shell_s1, shell_s2, shell_s3] + + result_no_screen = electron_repulsion_integral_improved(basis, notation="chemist") + result_screened = electron_repulsion_integral_improved( + basis, notation="chemist", schwarz_threshold=1e-12 + ) + + np.testing.assert_allclose(result_no_screen, result_screened, atol=1e-10) + + def test_schwarz_statistics(self): + """Verify screening statistics are tracked.""" + shell_s1 = GeneralizedContractionShell( + 0, np.array([0.0, 0.0, 0.0]), np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s2 = GeneralizedContractionShell( + 0, np.array([100.0, 0.0, 0.0]), np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + shell_s3 = GeneralizedContractionShell( + 0, np.array([0.0, 100.0, 0.0]), np.array([[1.0]]), np.array([1.0]), "cartesian" + ) + basis = [shell_s1, shell_s2, shell_s3] + + screener = SchwarzScreener( + list(basis), + PointChargeIntegral.boys_func, + compute_two_electron_integrals_os_hgp, + threshold=1e-10, + ) + stats = screener.get_statistics() + # Far-apart shells should have some bounds below threshold + assert stats["total"] == 0 # No quartets checked yet via is_significant + + # Now run full computation with screening + result = electron_repulsion_integral_improved( + basis, notation="chemist", schwarz_threshold=1e-10 + ) + assert np.all(np.isfinite(result)) From cdf9bd4e1e66308978bcc791ae4a8444afaa4eb0 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Sun, 8 Mar 2026 18:28:08 +0530 Subject: [PATCH 18/29] feat: add PySCF comparison tests and OS+HGP tutorial notebook --- ... Two-Electron Integral Optimizations.ipynb | 546 ++++++++++++++++++ tests/test_pyscf_comparison.py | 353 +++++++++++ 2 files changed, 899 insertions(+) create mode 100644 notebooks/tutorial/Tutorial 3 - Two-Electron Integral Optimizations.ipynb create mode 100644 tests/test_pyscf_comparison.py diff --git a/notebooks/tutorial/Tutorial 3 - Two-Electron Integral Optimizations.ipynb b/notebooks/tutorial/Tutorial 3 - Two-Electron Integral Optimizations.ipynb new file mode 100644 index 00000000..c49e25b9 --- /dev/null +++ b/notebooks/tutorial/Tutorial 3 - Two-Electron Integral Optimizations.ipynb @@ -0,0 +1,546 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# GBasis Tutorial - Two-Electron Integral Optimizations\n", + "\n", + "This tutorial covers the **optimized two-electron integral computation** in GBasis, including:\n", + "\n", + "1. **OS+HGP Algorithm** - Obara-Saika + Head-Gordon-Pople recursion scheme\n", + "2. **Boys Functions** - Standard Coulomb and range-separated (erf/erfc) potentials\n", + "3. **Schwarz Screening** - Skip negligible integrals for speedup\n", + "4. **High Angular Momentum** - Stable computation for d/f orbitals\n", + "5. **PySCF Comparison** - Validation against a trusted reference\n", + "\n", + "These optimizations are implemented as part of **Issue #221** to improve 2-electron integral performance." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Setup and Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import time\n", + "\n", + "# GBasis imports\n", + "from gbasis.contractions import GeneralizedContractionShell\n", + "from gbasis.integrals.electron_repulsion import (\n", + " ElectronRepulsionIntegral,\n", + " ElectronRepulsionIntegralImproved,\n", + " electron_repulsion_integral,\n", + " electron_repulsion_integral_improved,\n", + ")\n", + "\n", + "# Low-level OS+HGP algorithm\n", + "from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp\n", + "from gbasis.integrals.point_charge import PointChargeIntegral\n", + "\n", + "# Schwarz screening\n", + "from gbasis.integrals._screening import (\n", + " SchwarzScreener,\n", + " shell_pair_significant,\n", + " compute_schwarz_bounds,\n", + ")\n", + "\n", + "# Boys functions for different potentials\n", + "from gbasis.integrals.boys_functions import (\n", + " boys_function_standard,\n", + " boys_function_erf,\n", + " boys_function_erfc,\n", + " get_boys_function,\n", + ")\n", + "\n", + "print(\"All imports successful!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Creating Test Molecules\n", + "\n", + "We'll create hydrogen chain molecules to demonstrate the optimizations.\n", + "Larger separations between atoms lead to more sparsity and better screening." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def make_shell(coord, angmom, exps, coeffs):\n", + " \"\"\"Create a GeneralizedContractionShell.\"\"\"\n", + " return GeneralizedContractionShell(\n", + " angmom,\n", + " np.array(coord),\n", + " np.array(coeffs).reshape(-1, 1),\n", + " np.array(exps),\n", + " 'cartesian'\n", + " )\n", + "\n", + "def make_h_chain(n_atoms, spacing):\n", + " \"\"\"Create a hydrogen chain with STO-3G basis.\"\"\"\n", + " # STO-3G exponents and coefficients for H\n", + " exps = [3.42525091, 0.62391373, 0.16885540]\n", + " coeffs = [0.15432897, 0.53532814, 0.44463454]\n", + " \n", + " basis = []\n", + " for i in range(n_atoms):\n", + " coord = [i * spacing, 0.0, 0.0]\n", + " basis.append(make_shell(coord, 0, exps, coeffs))\n", + " \n", + " return basis\n", + "\n", + "# Create test molecules\n", + "h4_compact = make_h_chain(4, spacing=1.4) # Compact: H-H = 1.4 Bohr\n", + "h6_extended = make_h_chain(6, spacing=15.0) # Extended: H-H = 15 Bohr\n", + "\n", + "print(f\"Compact H4: {len(h4_compact)} shells, 1.4 Bohr spacing\")\n", + "print(f\"Extended H6: {len(h6_extended)} shells, 15.0 Bohr spacing\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. The OS+HGP Algorithm\n", + "\n", + "The improved implementation uses the **Obara-Saika + Head-Gordon-Pople** algorithm:\n", + "\n", + "```\n", + "Step 1: Boys Function F_m(T) = initial values\n", + "Step 2: VRR (Vertical) [a0|00]^m from [00|00]^m\n", + "Step 3: ETR (Electron) [a0|c0]^0 from [a0|00]^m \n", + "Step 4: Contract Sum over primitives (HGP: do early!)\n", + "Step 5: HRR (Horizontal) [ab|cd] from [a0|c0] (HGP: do last!)\n", + "```\n", + "\n", + "### Key Optimization (HGP):\n", + "The **Head-Gordon-Pople** optimization does HRR **after** contraction.\n", + "This is more efficient because contracted integrals are smaller than primitive integrals." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(\"\"\"\n", + "Obara-Saika + Head-Gordon-Pople Algorithm\n", + "==========================================\n", + "\n", + "Traditional: HGP Optimized:\n", + "------------ --------------\n", + "Boys F_m(T) Boys F_m(T)\n", + " | |\n", + " VRR VRR\n", + " | |\n", + " ETR ETR\n", + " | |\n", + " HRR <-- expensive! Contract <-- do this first!\n", + " | |\n", + " Contract HRR <-- now cheaper!\n", + " | |\n", + " [ab|cd] [ab|cd]\n", + "\n", + "Why is HGP better?\n", + "- HRR has NO 'm' index (no auxiliary function dependence)\n", + "- After contraction, fewer terms to process\n", + "- Especially beneficial for contracted basis sets (like cc-pVTZ)\n", + "\"\"\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Comparing Original vs Improved Implementation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Compare timing: original vs improved\n", + "basis = make_h_chain(4, spacing=1.4)\n", + "\n", + "# Original implementation\n", + "start = time.perf_counter()\n", + "eri_old = electron_repulsion_integral(basis, notation=\"chemist\")\n", + "time_old = time.perf_counter() - start\n", + "\n", + "# Improved OS+HGP implementation\n", + "start = time.perf_counter()\n", + "eri_new = electron_repulsion_integral_improved(basis, notation=\"chemist\")\n", + "time_new = time.perf_counter() - start\n", + "\n", + "# Verify they match\n", + "max_diff = np.max(np.abs(eri_old - eri_new))\n", + "\n", + "print(f\"Original implementation: {time_old:.4f} seconds\")\n", + "print(f\"Improved (OS+HGP): {time_new:.4f} seconds\")\n", + "print(f\"Max difference: {max_diff:.2e}\")\n", + "print(f\"Results match: {np.allclose(eri_old, eri_new, rtol=1e-10)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Boys Functions for Different Potentials\n", + "\n", + "GBasis supports various two-electron potentials beyond the standard Coulomb $1/r_{12}$:\n", + "\n", + "| Potential | Formula | Use Case |\n", + "|-----------|---------|----------|\n", + "| **Coulomb** | $1/r_{12}$ | Standard electron repulsion |\n", + "| **erf (long-range)** | $\\text{erf}(\\omega r)/r$ | Range-separated hybrids (LC-) |\n", + "| **erfc (short-range)** | $\\text{erfc}(\\omega r)/r$ | Short-range exchange |\n", + "\n", + "The key identity: $\\text{erf}(\\omega r)/r + \\text{erfc}(\\omega r)/r = 1/r$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Compare Boys functions for different potentials\n", + "T_values = np.array([0.0, 0.5, 1.0, 2.0, 5.0])\n", + "m = np.array([0]) # F_0(T)\n", + "omega = 0.5 # Range-separation parameter\n", + "\n", + "print(\"Boys Function F_0(T) for Different Potentials:\")\n", + "print(\"-\" * 50)\n", + "print(f\"{'T':>8} {'Coulomb':>12} {'erf':>12} {'erfc':>12}\")\n", + "print(\"-\" * 50)\n", + "\n", + "for t in T_values:\n", + " t_arr = np.array([t])\n", + " rho = np.array([0.8])\n", + " \n", + " coulomb = boys_function_standard(m, t_arr)[0]\n", + " erf_val = boys_function_erf(m, t_arr, rho, omega)[0]\n", + " erfc_val = boys_function_erfc(m, t_arr, rho, omega)[0]\n", + " \n", + " print(f\"{t:>8.1f} {coulomb:>12.6f} {erf_val:>12.6f} {erfc_val:>12.6f}\")\n", + "\n", + "print(\"\\nNote: erf + erfc = Coulomb (range separation identity)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify: erf + erfc = Coulomb\n", + "T_test = np.array([1.0])\n", + "rho = np.array([0.8])\n", + "\n", + "coulomb = boys_function_standard(m, T_test)[0]\n", + "erf_val = boys_function_erf(m, T_test, rho, omega)[0]\n", + "erfc_val = boys_function_erfc(m, T_test, rho, omega)[0]\n", + "\n", + "print(\"Range Separation Verification:\")\n", + "print(f\" Coulomb: {coulomb:.10f}\")\n", + "print(f\" erf + erfc: {erf_val + erfc_val:.10f}\")\n", + "print(f\" Difference: {abs(coulomb - (erf_val + erfc_val)):.2e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Using the `get_boys_function` Factory" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Factory function for different potentials\n", + "print(\"Using get_boys_function() factory:\")\n", + "print(\"-\" * 40)\n", + "\n", + "# Standard Coulomb\n", + "boys_coulomb = get_boys_function(\"coulomb\")\n", + "print(f\"Coulomb: F_0(1.0) = {boys_coulomb(np.array([0]), np.array([1.0]))[0]:.6f}\")\n", + "\n", + "# Long-range (erf)\n", + "boys_lr = get_boys_function(\"erf\", omega=0.5)\n", + "print(f\"Long-range (omega=0.5): F_0(1.0) = {boys_lr(np.array([0]), np.array([1.0]), rho=np.array([0.8]))[0]:.6f}\")\n", + "\n", + "# Short-range (erfc)\n", + "boys_sr = get_boys_function(\"erfc\", omega=0.5)\n", + "print(f\"Short-range (omega=0.5): F_0(1.0) = {boys_sr(np.array([0]), np.array([1.0]), rho=np.array([0.8]))[0]:.6f}\")\n", + "\n", + "print(\"\\nAvailable potentials: 'coulomb', 'erf', 'erfc'\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Schwarz Screening\n", + "\n", + "### What is Schwarz Screening?\n", + "\n", + "The **Schwarz inequality** states:\n", + "\n", + "$$|(ab|cd)| \\leq \\sqrt{(ab|ab)} \\times \\sqrt{(cd|cd)}$$\n", + "\n", + "If this upper bound is smaller than a threshold (e.g., $10^{-12}$), we can skip the integral computation entirely!\n", + "\n", + "### Benefits:\n", + "- For extended molecules, 60-90% of integrals can be skipped\n", + "- Speedup of 2-10x depending on system size and geometry" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Demonstrate Schwarz screening on extended system\n", + "boys_func = PointChargeIntegral.boys_func\n", + "\n", + "print(\"=\" * 55)\n", + "print(\"Schwarz Screening Comparison\")\n", + "print(\"=\" * 55)\n", + "\n", + "for name, basis, spacing in [\n", + " (\"Compact H4 (1.4 Bohr)\", h4_compact, 1.4),\n", + " (\"Extended H6 (15.0 Bohr)\", h6_extended, 15.0),\n", + "]:\n", + " screener = SchwarzScreener(\n", + " basis, boys_func, compute_two_electron_integrals_os_hgp\n", + " )\n", + " \n", + " n = len(basis)\n", + " for i in range(n):\n", + " for j in range(n):\n", + " for k in range(n):\n", + " for l in range(n):\n", + " screener.is_significant(i, j, k, l)\n", + " \n", + " stats = screener.get_statistics()\n", + " \n", + " print(f\"\\n{name}:\")\n", + " print(f\" Total shell quartets: {stats['total']}\")\n", + " print(f\" Computed: {stats['n_computed']}\")\n", + " print(f\" Screened (skipped): {stats['n_screened']}\")\n", + " print(f\" Percent screened: {stats['percent_screened']:.1f}%\")\n", + " print(f\" Speedup factor: {stats['speedup_factor']:.2f}x\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Shell-Pair Prescreening\n", + "\n", + "Before even computing Schwarz bounds, we can use the **Gaussian product theorem** to\n", + "quickly check if a shell pair is significant:\n", + "\n", + "$$\\exp\\left(-\\frac{\\alpha_a \\alpha_b}{\\alpha_a + \\alpha_b} |\\mathbf{A} - \\mathbf{B}|^2\\right)$$\n", + "\n", + "If this decay factor is below threshold for all primitive pairs, the shell pair contributes nothing." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Demonstrate shell-pair prescreening\n", + "print(\"Shell-Pair Prescreening:\")\n", + "print(\"-\" * 50)\n", + "\n", + "# Same center - always significant\n", + "shell_a = make_shell([0, 0, 0], 0, [1.0], [1.0])\n", + "shell_b = make_shell([0, 0, 0], 0, [0.5], [1.0])\n", + "print(f\"Same center (0,0,0): {shell_pair_significant(shell_a, shell_b)}\")\n", + "\n", + "# Close shells\n", + "shell_c = make_shell([1.0, 0, 0], 0, [1.0], [1.0])\n", + "print(f\"Close (1.0 Bohr apart): {shell_pair_significant(shell_a, shell_c)}\")\n", + "\n", + "# Far apart\n", + "shell_d = make_shell([50.0, 0, 0], 0, [1.0], [1.0])\n", + "print(f\"Far apart (50 Bohr): {shell_pair_significant(shell_a, shell_d)}\")\n", + "\n", + "# Very diffuse (small exponent) at distance\n", + "shell_e = make_shell([10.0, 0, 0], 0, [0.01], [1.0])\n", + "print(f\"Diffuse at 10 Bohr (a=0.01): {shell_pair_significant(shell_a, shell_e)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. High Angular Momentum Integrals\n", + "\n", + "The improved algorithm correctly handles **d-orbitals (L=2)** and **f-orbitals (L=3)** without numerical overflow." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "coord = np.array([0.0, 0.0, 0.0])\n", + "exps = np.array([0.5])\n", + "coeffs = np.array([[1.0]])\n", + "\n", + "# Angular momentum components\n", + "orbitals = {\n", + " 's (L=0)': (0, np.array([[0, 0, 0]])),\n", + " 'p (L=1)': (1, np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])),\n", + " 'd (L=2)': (2, np.array([[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0], [0, 1, 1], [0, 0, 2]])),\n", + " 'f (L=3)': (3, np.array([[3, 0, 0], [2, 1, 0], [2, 0, 1], [1, 2, 0], [1, 1, 1],\n", + " [1, 0, 2], [0, 3, 0], [0, 2, 1], [0, 1, 2], [0, 0, 3]]))\n", + "}\n", + "\n", + "print(\"High Angular Momentum Test:\")\n", + "print(\"-\" * 60)\n", + "print(f\"{'Orbital':>10} {'Shape':>20} {'Has NaN':>10} {'Has Inf':>10}\")\n", + "print(\"-\" * 60)\n", + "\n", + "for name, (angmom, comp) in orbitals.items():\n", + " result = compute_two_electron_integrals_os_hgp(\n", + " PointChargeIntegral.boys_func,\n", + " coord, angmom, comp, exps, coeffs,\n", + " coord, angmom, comp, exps, coeffs,\n", + " coord, angmom, comp, exps, coeffs,\n", + " coord, angmom, comp, exps, coeffs,\n", + " )\n", + " \n", + " shape_str = str(result.shape[:4])\n", + " has_nan = np.any(np.isnan(result))\n", + " has_inf = np.any(np.isinf(result))\n", + " \n", + " print(f\"{name:>10} {shape_str:>20} {str(has_nan):>10} {str(has_inf):>10}\")\n", + "\n", + "print(\"\\nAll integrals computed successfully without numerical issues!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Validation Against PySCF\n", + "\n", + "We validate our implementation against PySCF, a trusted quantum chemistry package.\n", + "Both the original and improved implementations should match PySCF exactly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " from pyscf import gto\n", + " HAS_PYSCF = True\n", + "except ImportError:\n", + " HAS_PYSCF = False\n", + " print(\"PySCF not installed. Install with: pip install pyscf\")\n", + "\n", + "if HAS_PYSCF:\n", + " # H2 with STO-3G basis\n", + " exps_h = np.array([3.42525091, 0.62391373, 0.16885540])\n", + " coeffs_h = np.array([[0.15432897], [0.53532814], [0.44463454]])\n", + " \n", + " basis_gbasis = [\n", + " GeneralizedContractionShell(0, np.array([0.0, 0.0, 0.0]), coeffs_h, exps_h, 'cartesian'),\n", + " GeneralizedContractionShell(0, np.array([1.4, 0.0, 0.0]), coeffs_h, exps_h, 'cartesian'),\n", + " ]\n", + " \n", + " mol = gto.M(\n", + " atom='H 0 0 0; H 1.4 0 0', unit='bohr',\n", + " basis={'H': gto.basis.parse(\n", + " 'H S\\n 3.42525091 0.15432897\\n 0.62391373 0.53532814\\n 0.16885540 0.44463454'\n", + " )}, verbose=0,\n", + " )\n", + " \n", + " # PySCF returns chemist notation (ij|kl)\n", + " eri_pyscf = mol.intor('int2e')\n", + " \n", + " # GBasis - both implementations in chemist notation\n", + " eri_original = electron_repulsion_integral(basis_gbasis, notation='chemist')\n", + " eri_improved = electron_repulsion_integral_improved(basis_gbasis, notation='chemist')\n", + " \n", + " print(\"H2/STO-3G Comparison (Chemist notation):\")\n", + " print(\"=\" * 55)\n", + " print(f\"{'Integral':>12} {'PySCF':>12} {'Original':>12} {'Improved':>12}\")\n", + " print(\"-\" * 55)\n", + " \n", + " for i in range(2):\n", + " for j in range(2):\n", + " label = f\"({i}{i}|{j}{j})\"\n", + " print(f\"{label:>12} {eri_pyscf[i,i,j,j]:>12.6f} {eri_original[i,i,j,j]:>12.6f} {eri_improved[i,i,j,j]:>12.6f}\")\n", + " \n", + " print(f\"\\nMax |Original - PySCF|: {np.max(np.abs(eri_original - eri_pyscf)):.2e}\")\n", + " print(f\"Max |Improved - PySCF|: {np.max(np.abs(eri_improved - eri_pyscf)):.2e}\")\n", + " print(f\"Max |Improved - Original|: {np.max(np.abs(eri_improved - eri_original)):.2e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 8. Summary\n\n### Key Features Implemented (Issue #221):\n\n| Week | Feature | Benefit |\n|------|---------|--------|\n| 1 | Boys Function | Efficient $F_m(T)$ via hyp1f1 for all orders |\n| 2 | VRR (Vertical Recursion) | Build angular momentum on center A |\n| 3 | ETR + Contraction | Transfer to center C + sum over primitives |\n| 4 | HRR (Horizontal Recursion) | Distribute to B and D (post-contraction per HGP) |\n| 5 | erf/erfc Boys Functions | Range-separated DFT support |\n| 6 | GBasis Integration | `ElectronRepulsionIntegralImproved` class |\n| 7 | Schwarz Screening | Skip negligible integrals (2-10x speedup) |\n| 8 | PySCF Validation | Verified against trusted reference |\n\n### When to Use Screening:\n- **Large molecules** with spatially separated atoms\n- **Extended basis sets** (cc-pVDZ, cc-pVTZ, etc.)\n- **Periodic systems** or **clusters**\n\n### When Screening Has Less Effect:\n- **Compact molecules** where all atoms are close\n- **Minimal basis sets** (STO-3G) with few primitives" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(\"\"\"\n", + "Thanks for completing this tutorial!\n", + "\n", + "For more information:\n", + "- Documentation: http://gbasis.qcdevs.org/\n", + "- GitHub: https://github.com/theochem/gbasis\n", + "- Issue #221: Two-electron integral improvements\n", + "\"\"\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/tests/test_pyscf_comparison.py b/tests/test_pyscf_comparison.py new file mode 100644 index 00000000..a490d730 --- /dev/null +++ b/tests/test_pyscf_comparison.py @@ -0,0 +1,353 @@ +"""Direct comparison of GBasis ERIs with PySCF reference. + +This test validates that both the original and improved (OS+HGP) GBasis +implementations produce the same electron repulsion integrals as PySCF. + +Notation conventions: +- PySCF mol.intor('int2e') returns integrals in chemist's notation (ij|kl). +- GBasis electron_repulsion_integral() defaults to physicist's notation . +- GBasis electron_repulsion_integral(notation="chemist") returns chemist's (ij|kl). +- Conversion: = (ik|jl), so physicist = chemist.transpose(0, 2, 1, 3). + +We use manual basis definitions to ensure both codes use exactly the same +basis parameters, avoiding normalization conversion issues. +""" + +import numpy as np +import pytest + +try: + from pyscf import gto + HAS_PYSCF = True +except ImportError: + HAS_PYSCF = False + +from gbasis.contractions import GeneralizedContractionShell +from gbasis.integrals.electron_repulsion import ( + electron_repulsion_integral, + electron_repulsion_integral_improved, +) + + +# --------------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------------- + +def make_h2_sto3g(): + """Create H2 with STO-3G basis (same parameters for GBasis and PySCF). + + STO-3G for H: 3 primitives contracted to 1 s-function. + Returns GBasis basis list and PySCF mol object. + """ + exps = np.array([3.42525091, 0.62391373, 0.16885540]) + coeffs = np.array([[0.15432897], [0.53532814], [0.44463454]]) + + coord1 = np.array([0.0, 0.0, 0.0]) + coord2 = np.array([1.4, 0.0, 0.0]) # 1.4 Bohr apart + + basis = [ + GeneralizedContractionShell(0, coord1, coeffs, exps, 'cartesian'), + GeneralizedContractionShell(0, coord2, coeffs, exps, 'cartesian'), + ] + + mol = gto.M( + atom='H 0 0 0; H 1.4 0 0', + unit='bohr', + basis={'H': gto.basis.parse(''' + H S + 3.42525091 0.15432897 + 0.62391373 0.53532814 + 0.16885540 0.44463454 + ''')}, + verbose=0, + ) + return basis, mol + + +def make_h2_primitive(): + """Create H2 with a single primitive s-function per atom. + + Simplest possible case: no contraction complexity. + """ + exps = np.array([1.0]) + coeffs = np.array([[1.0]]) + + coord1 = np.array([0.0, 0.0, 0.0]) + coord2 = np.array([2.0, 0.0, 0.0]) + + basis = [ + GeneralizedContractionShell(0, coord1, coeffs, exps, 'cartesian'), + GeneralizedContractionShell(0, coord2, coeffs, exps, 'cartesian'), + ] + + mol = gto.M( + atom='H 0 0 0; H 2.0 0 0', + unit='bohr', + basis={'H': gto.basis.parse('H S\n 1.0 1.0')}, + verbose=0, + ) + return basis, mol + + +def make_h2_sp(): + """Create H2 with s + p basis (one s and one p per atom). + + This tests angular momentum handling, which is critical because + chemist/physicist notation differences matter for mixed angular momentum. + """ + s_exps = np.array([1.0]) + s_coeffs = np.array([[1.0]]) + p_exps = np.array([0.8]) + p_coeffs = np.array([[1.0]]) + + coord1 = np.array([0.0, 0.0, 0.0]) + coord2 = np.array([2.0, 0.0, 0.0]) + + basis = [ + GeneralizedContractionShell(0, coord1, s_coeffs, s_exps, 'cartesian'), + GeneralizedContractionShell(1, coord1, p_coeffs, p_exps, 'cartesian'), + GeneralizedContractionShell(0, coord2, s_coeffs, s_exps, 'cartesian'), + GeneralizedContractionShell(1, coord2, p_coeffs, p_exps, 'cartesian'), + ] + + mol = gto.M( + atom='H 0 0 0; H 2.0 0 0', + unit='bohr', + basis={'H': gto.basis.parse(''' + H S + 1.0 1.0 + H P + 0.8 1.0 + ''')}, + verbose=0, + ) + return basis, mol + + +# --------------------------------------------------------------------------- +# Tests: Original implementation vs PySCF +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not HAS_PYSCF, reason="PySCF not installed") +class TestPySCFOriginal: + """Compare original GBasis ERIs with PySCF reference.""" + + def test_h2_sto3g_chemist(self): + """Test H2/STO-3G in chemist notation: GBasis (ij|kl) == PySCF (ij|kl).""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') # PySCF returns chemist (ij|kl) + + np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, + err_msg="H2/STO-3G chemist notation doesn't match PySCF") + + def test_h2_sto3g_physicist(self): + """Test H2/STO-3G in physicist notation. + + GBasis physicist == PySCF chemist (ij|kl).transpose(0,2,1,3). + """ + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="physicist") + eri_pyscf = mol.intor('int2e') + eri_pyscf_physicist = eri_pyscf.transpose(0, 2, 1, 3) + + np.testing.assert_allclose(eri_gbasis, eri_pyscf_physicist, rtol=1e-6, atol=1e-10, + err_msg="H2/STO-3G physicist notation doesn't match PySCF") + + def test_h2_primitive(self): + """Test single-primitive H2 (simplest possible case).""" + basis, mol = make_h2_primitive() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') + + np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, + err_msg="Single primitive H2 ERIs don't match PySCF") + + def test_h2_sp_basis(self): + """Test H2 with s+p basis (tests angular momentum handling).""" + basis, mol = make_h2_sp() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') + + np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, + err_msg="H2 s+p basis ERIs don't match PySCF") + + def test_h2_specific_integrals(self): + """Test specific physically meaningful integrals for H2/STO-3G.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') + + # (00|00) - Coulomb integral, both electrons on atom 1 + assert np.isclose(eri_gbasis[0, 0, 0, 0], eri_pyscf[0, 0, 0, 0], rtol=1e-6), \ + f"(00|00) mismatch: GBasis={eri_gbasis[0,0,0,0]:.10f}, PySCF={eri_pyscf[0,0,0,0]:.10f}" + + # (11|11) - Coulomb integral, both electrons on atom 2 + assert np.isclose(eri_gbasis[1, 1, 1, 1], eri_pyscf[1, 1, 1, 1], rtol=1e-6), \ + f"(11|11) mismatch: GBasis={eri_gbasis[1,1,1,1]:.10f}, PySCF={eri_pyscf[1,1,1,1]:.10f}" + + # By symmetry, (00|00) == (11|11) for identical atoms + assert np.isclose(eri_gbasis[0, 0, 0, 0], eri_gbasis[1, 1, 1, 1], rtol=1e-6), \ + "Identical atom diagonal integrals should be equal" + + # (01|01) - exchange-type integral + assert np.isclose(eri_gbasis[0, 1, 0, 1], eri_pyscf[0, 1, 0, 1], rtol=1e-6), \ + f"(01|01) mismatch" + + def test_h2_symmetries(self): + """Test that PySCF and GBasis agree on 8-fold ERI symmetries.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral(basis, notation="chemist") + + # Chemist notation symmetries: (ij|kl) = (ji|kl) = (ij|lk) = (kl|ij) + n = eri_gbasis.shape[0] + for i in range(n): + for j in range(n): + for k in range(n): + for l in range(n): + val = eri_gbasis[i, j, k, l] + assert np.isclose(val, eri_gbasis[j, i, k, l], rtol=1e-10), \ + f"(ij|kl)!=(ji|kl) for ({i}{j}|{k}{l})" + assert np.isclose(val, eri_gbasis[i, j, l, k], rtol=1e-10), \ + f"(ij|kl)!=(ij|lk) for ({i}{j}|{k}{l})" + assert np.isclose(val, eri_gbasis[k, l, i, j], rtol=1e-10), \ + f"(ij|kl)!=(kl|ij) for ({i}{j}|{k}{l})" + + +# --------------------------------------------------------------------------- +# Tests: Improved (OS+HGP) implementation vs PySCF +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not HAS_PYSCF, reason="PySCF not installed") +class TestPySCFImproved: + """Compare improved OS+HGP GBasis ERIs with PySCF reference.""" + + def test_h2_sto3g_chemist(self): + """Test improved implementation H2/STO-3G in chemist notation.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') + + np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, + err_msg="Improved H2/STO-3G chemist notation doesn't match PySCF") + + def test_h2_sto3g_physicist(self): + """Test improved implementation H2/STO-3G in physicist notation.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="physicist") + eri_pyscf = mol.intor('int2e') + eri_pyscf_physicist = eri_pyscf.transpose(0, 2, 1, 3) + + np.testing.assert_allclose(eri_gbasis, eri_pyscf_physicist, rtol=1e-6, atol=1e-10, + err_msg="Improved H2/STO-3G physicist notation doesn't match PySCF") + + def test_h2_primitive(self): + """Test improved implementation with single primitive.""" + basis, mol = make_h2_primitive() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') + + np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, + err_msg="Improved single primitive H2 ERIs don't match PySCF") + + def test_h2_sp_basis(self): + """Test improved implementation H2 with s+p basis.""" + basis, mol = make_h2_sp() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') + + np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, + err_msg="Improved H2 s+p basis ERIs don't match PySCF") + + def test_improved_matches_original(self): + """Test that improved and original implementations agree on PySCF test case.""" + basis, _ = make_h2_sto3g() + + eri_old = electron_repulsion_integral(basis, notation="chemist") + eri_new = electron_repulsion_integral_improved(basis, notation="chemist") + + np.testing.assert_allclose(eri_new, eri_old, rtol=1e-10, + err_msg="Improved doesn't match original for H2/STO-3G") + + def test_improved_matches_original_sp(self): + """Test that improved and original agree for s+p basis.""" + basis, _ = make_h2_sp() + + eri_old = electron_repulsion_integral(basis, notation="chemist") + eri_new = electron_repulsion_integral_improved(basis, notation="chemist") + + np.testing.assert_allclose(eri_new, eri_old, rtol=1e-10, + err_msg="Improved doesn't match original for H2 s+p") + + +# --------------------------------------------------------------------------- +# Tests: Physical sanity checks +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not HAS_PYSCF, reason="PySCF not installed") +class TestPySCFPhysicalProperties: + """Verify physical properties of ERIs using PySCF as cross-check.""" + + def test_positive_diagonal(self): + """Test that Coulomb integrals (ii|ii) are positive.""" + basis, mol = make_h2_sto3g() + + eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") + eri_pyscf = mol.intor('int2e') + + n = eri_gbasis.shape[0] + for i in range(n): + assert eri_gbasis[i, i, i, i] > 0, f"GBasis (ii|ii) not positive for i={i}" + assert eri_pyscf[i, i, i, i] > 0, f"PySCF (ii|ii) not positive for i={i}" + + def test_coulomb_greater_than_exchange(self): + """Test that Coulomb integral >= exchange integral for H2. + + For same-atom shells: (00|00) >= (01|01) because the exchange + integral involves orbital overlap which reduces the value. + """ + basis, mol = make_h2_sto3g() + + eri = electron_repulsion_integral_improved(basis, notation="chemist") + + # Coulomb (00|00) should be > exchange (01|01) + coulomb = eri[0, 0, 0, 0] + exchange = eri[0, 1, 0, 1] + assert coulomb > exchange, \ + f"Coulomb ({coulomb:.6f}) should be > exchange ({exchange:.6f})" + + def test_eri_values_physically_reasonable(self): + """Test that ERI values are in a physically reasonable range. + + For H2 with STO-3G, typical values are 0.1-1.0 in atomic units. + """ + basis, _ = make_h2_sto3g() + + eri = electron_repulsion_integral_improved(basis, notation="chemist") + + # All values should be finite + assert np.all(np.isfinite(eri)), "ERIs contain NaN or Inf" + + # Diagonal integrals should be in reasonable range (0.1 - 2.0 a.u.) + for i in range(eri.shape[0]): + val = eri[i, i, i, i] + assert 0.01 < val < 10.0, \ + f"Diagonal integral ({i}{i}|{i}{i}) = {val} outside reasonable range" + + +if __name__ == '__main__': + if not HAS_PYSCF: + print("PySCF not installed. Install with: pip install pyscf") + print("Skipping PySCF comparison tests.") + else: + pytest.main([__file__, '-v']) From f630272e6983c497e24090e3e559890a57cf6a86 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Sun, 8 Mar 2026 21:37:27 +0530 Subject: [PATCH 19/29] fix: tighten libcint ERI tolerance and use improved implementation (#208) --- tests/test_libcint.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/test_libcint.py b/tests/test_libcint.py index 51c12667..1d1deddb 100644 --- a/tests/test_libcint.py +++ b/tests/test_libcint.py @@ -13,7 +13,10 @@ import gbasis from gbasis.integrals.angular_momentum import angular_momentum_integral -from gbasis.integrals.electron_repulsion import electron_repulsion_integral +from gbasis.integrals.electron_repulsion import ( + electron_repulsion_integral, + electron_repulsion_integral_improved, +) from gbasis.integrals.kinetic_energy import kinetic_energy_integral from gbasis.integrals.moment import moment_integral from gbasis.integrals.momentum import momentum_integral @@ -80,7 +83,7 @@ def test_integral(basis, atsyms, atcoords, coord_type, integral): against the GBasis Python integrals. """ - atol, rtol = 1e-4, 1e-4 + atol, rtol = 1e-6, 1e-6 atcoords = atcoords / 0.5291772083 @@ -125,7 +128,7 @@ def test_integral(basis, atsyms, atcoords, coord_type, integral): npt.assert_array_equal(lc_int.shape, (lc_basis.nbfn, lc_basis.nbfn, 3)) elif integral == "electron_repulsion": - py_int = electron_repulsion_integral(py_basis) + py_int = electron_repulsion_integral_improved(py_basis) npt.assert_array_equal( py_int.shape, (lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn) ) @@ -133,6 +136,8 @@ def test_integral(basis, atsyms, atcoords, coord_type, integral): npt.assert_array_equal( lc_int.shape, (lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn, lc_basis.nbfn) ) + npt.assert_allclose(lc_int, py_int, atol=1e-4, rtol=1e-5) + return elif integral == "point_charge": charge_coords = np.asarray([[2.0, 2.0, 2.0], [-3.0, -3.0, -3.0], [-1.0, 2.0, -3.0]]) @@ -204,7 +209,7 @@ def test_integral_iodata(fname, elements, coord_type, integral, transform): from iodata import load_one from gbasis.integrals.libcint import ELEMENTS, LIBCINT, CBasis - atol, rtol = 1e-4, 1e-4 + atol, rtol = 1e-6, 1e-6 mol=load_one(find_datafile(fname)) py_basis=from_iodata(mol) From b351377ff8452b9e6c0b67aa1eb8f4a545f7e398 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Wed, 11 Mar 2026 22:09:52 +0530 Subject: [PATCH 20/29] fix: add pyscf to dev deps and remove conditional imports in test --- pyproject.toml | 1 + tests/test_pyscf_comparison.py | 214 +++++++++++++++++++-------------- 2 files changed, 128 insertions(+), 87 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 82727343..3dfa1618 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dev = [ "sphinx", "sphinx_autodoc_typehints", "sphinx-copybutton", + "pyscf>=1.6.1", ] doc = [ "numpydoc", diff --git a/tests/test_pyscf_comparison.py b/tests/test_pyscf_comparison.py index a490d730..78965159 100644 --- a/tests/test_pyscf_comparison.py +++ b/tests/test_pyscf_comparison.py @@ -15,12 +15,7 @@ import numpy as np import pytest - -try: - from pyscf import gto - HAS_PYSCF = True -except ImportError: - HAS_PYSCF = False +from pyscf import gto from gbasis.contractions import GeneralizedContractionShell from gbasis.integrals.electron_repulsion import ( @@ -28,11 +23,11 @@ electron_repulsion_integral_improved, ) - # --------------------------------------------------------------------------- # Helper functions # --------------------------------------------------------------------------- + def make_h2_sto3g(): """Create H2 with STO-3G basis (same parameters for GBasis and PySCF). @@ -46,19 +41,19 @@ def make_h2_sto3g(): coord2 = np.array([1.4, 0.0, 0.0]) # 1.4 Bohr apart basis = [ - GeneralizedContractionShell(0, coord1, coeffs, exps, 'cartesian'), - GeneralizedContractionShell(0, coord2, coeffs, exps, 'cartesian'), + GeneralizedContractionShell(0, coord1, coeffs, exps, "cartesian"), + GeneralizedContractionShell(0, coord2, coeffs, exps, "cartesian"), ] mol = gto.M( - atom='H 0 0 0; H 1.4 0 0', - unit='bohr', - basis={'H': gto.basis.parse(''' + atom="H 0 0 0; H 1.4 0 0", + unit="bohr", + basis={"H": gto.basis.parse(""" H S 3.42525091 0.15432897 0.62391373 0.53532814 0.16885540 0.44463454 - ''')}, + """)}, verbose=0, ) return basis, mol @@ -76,14 +71,14 @@ def make_h2_primitive(): coord2 = np.array([2.0, 0.0, 0.0]) basis = [ - GeneralizedContractionShell(0, coord1, coeffs, exps, 'cartesian'), - GeneralizedContractionShell(0, coord2, coeffs, exps, 'cartesian'), + GeneralizedContractionShell(0, coord1, coeffs, exps, "cartesian"), + GeneralizedContractionShell(0, coord2, coeffs, exps, "cartesian"), ] mol = gto.M( - atom='H 0 0 0; H 2.0 0 0', - unit='bohr', - basis={'H': gto.basis.parse('H S\n 1.0 1.0')}, + atom="H 0 0 0; H 2.0 0 0", + unit="bohr", + basis={"H": gto.basis.parse("H S\n 1.0 1.0")}, verbose=0, ) return basis, mol @@ -104,21 +99,21 @@ def make_h2_sp(): coord2 = np.array([2.0, 0.0, 0.0]) basis = [ - GeneralizedContractionShell(0, coord1, s_coeffs, s_exps, 'cartesian'), - GeneralizedContractionShell(1, coord1, p_coeffs, p_exps, 'cartesian'), - GeneralizedContractionShell(0, coord2, s_coeffs, s_exps, 'cartesian'), - GeneralizedContractionShell(1, coord2, p_coeffs, p_exps, 'cartesian'), + GeneralizedContractionShell(0, coord1, s_coeffs, s_exps, "cartesian"), + GeneralizedContractionShell(1, coord1, p_coeffs, p_exps, "cartesian"), + GeneralizedContractionShell(0, coord2, s_coeffs, s_exps, "cartesian"), + GeneralizedContractionShell(1, coord2, p_coeffs, p_exps, "cartesian"), ] mol = gto.M( - atom='H 0 0 0; H 2.0 0 0', - unit='bohr', - basis={'H': gto.basis.parse(''' + atom="H 0 0 0; H 2.0 0 0", + unit="bohr", + basis={"H": gto.basis.parse(""" H S 1.0 1.0 H P 0.8 1.0 - ''')}, + """)}, verbose=0, ) return basis, mol @@ -128,7 +123,7 @@ def make_h2_sp(): # Tests: Original implementation vs PySCF # --------------------------------------------------------------------------- -@pytest.mark.skipif(not HAS_PYSCF, reason="PySCF not installed") + class TestPySCFOriginal: """Compare original GBasis ERIs with PySCF reference.""" @@ -137,10 +132,15 @@ def test_h2_sto3g_chemist(self): basis, mol = make_h2_sto3g() eri_gbasis = electron_repulsion_integral(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') # PySCF returns chemist (ij|kl) + eri_pyscf = mol.intor("int2e") # PySCF returns chemist (ij|kl) - np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, - err_msg="H2/STO-3G chemist notation doesn't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="H2/STO-3G chemist notation doesn't match PySCF", + ) def test_h2_sto3g_physicist(self): """Test H2/STO-3G in physicist notation. @@ -150,58 +150,77 @@ def test_h2_sto3g_physicist(self): basis, mol = make_h2_sto3g() eri_gbasis = electron_repulsion_integral(basis, notation="physicist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") eri_pyscf_physicist = eri_pyscf.transpose(0, 2, 1, 3) - np.testing.assert_allclose(eri_gbasis, eri_pyscf_physicist, rtol=1e-6, atol=1e-10, - err_msg="H2/STO-3G physicist notation doesn't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf_physicist, + rtol=1e-6, + atol=1e-10, + err_msg="H2/STO-3G physicist notation doesn't match PySCF", + ) def test_h2_primitive(self): """Test single-primitive H2 (simplest possible case).""" basis, mol = make_h2_primitive() eri_gbasis = electron_repulsion_integral(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") - np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, - err_msg="Single primitive H2 ERIs don't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Single primitive H2 ERIs don't match PySCF", + ) def test_h2_sp_basis(self): """Test H2 with s+p basis (tests angular momentum handling).""" basis, mol = make_h2_sp() eri_gbasis = electron_repulsion_integral(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") - np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, - err_msg="H2 s+p basis ERIs don't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="H2 s+p basis ERIs don't match PySCF", + ) def test_h2_specific_integrals(self): """Test specific physically meaningful integrals for H2/STO-3G.""" basis, mol = make_h2_sto3g() eri_gbasis = electron_repulsion_integral(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") # (00|00) - Coulomb integral, both electrons on atom 1 - assert np.isclose(eri_gbasis[0, 0, 0, 0], eri_pyscf[0, 0, 0, 0], rtol=1e-6), \ - f"(00|00) mismatch: GBasis={eri_gbasis[0,0,0,0]:.10f}, PySCF={eri_pyscf[0,0,0,0]:.10f}" + assert np.isclose( + eri_gbasis[0, 0, 0, 0], eri_pyscf[0, 0, 0, 0], rtol=1e-6 + ), f"(00|00) mismatch: GBasis={eri_gbasis[0,0,0,0]:.10f}, PySCF={eri_pyscf[0,0,0,0]:.10f}" # (11|11) - Coulomb integral, both electrons on atom 2 - assert np.isclose(eri_gbasis[1, 1, 1, 1], eri_pyscf[1, 1, 1, 1], rtol=1e-6), \ - f"(11|11) mismatch: GBasis={eri_gbasis[1,1,1,1]:.10f}, PySCF={eri_pyscf[1,1,1,1]:.10f}" + assert np.isclose( + eri_gbasis[1, 1, 1, 1], eri_pyscf[1, 1, 1, 1], rtol=1e-6 + ), f"(11|11) mismatch: GBasis={eri_gbasis[1,1,1,1]:.10f}, PySCF={eri_pyscf[1,1,1,1]:.10f}" # By symmetry, (00|00) == (11|11) for identical atoms - assert np.isclose(eri_gbasis[0, 0, 0, 0], eri_gbasis[1, 1, 1, 1], rtol=1e-6), \ - "Identical atom diagonal integrals should be equal" + assert np.isclose( + eri_gbasis[0, 0, 0, 0], eri_gbasis[1, 1, 1, 1], rtol=1e-6 + ), "Identical atom diagonal integrals should be equal" # (01|01) - exchange-type integral - assert np.isclose(eri_gbasis[0, 1, 0, 1], eri_pyscf[0, 1, 0, 1], rtol=1e-6), \ - f"(01|01) mismatch" + assert np.isclose( + eri_gbasis[0, 1, 0, 1], eri_pyscf[0, 1, 0, 1], rtol=1e-6 + ), "(01|01) mismatch" def test_h2_symmetries(self): """Test that PySCF and GBasis agree on 8-fold ERI symmetries.""" - basis, mol = make_h2_sto3g() + basis, _mol = make_h2_sto3g() eri_gbasis = electron_repulsion_integral(basis, notation="chemist") @@ -210,21 +229,24 @@ def test_h2_symmetries(self): for i in range(n): for j in range(n): for k in range(n): - for l in range(n): - val = eri_gbasis[i, j, k, l] - assert np.isclose(val, eri_gbasis[j, i, k, l], rtol=1e-10), \ - f"(ij|kl)!=(ji|kl) for ({i}{j}|{k}{l})" - assert np.isclose(val, eri_gbasis[i, j, l, k], rtol=1e-10), \ - f"(ij|kl)!=(ij|lk) for ({i}{j}|{k}{l})" - assert np.isclose(val, eri_gbasis[k, l, i, j], rtol=1e-10), \ - f"(ij|kl)!=(kl|ij) for ({i}{j}|{k}{l})" + for m in range(n): + val = eri_gbasis[i, j, k, m] + assert np.isclose( + val, eri_gbasis[j, i, k, m], rtol=1e-10 + ), f"(ij|kl)!=(ji|kl) for ({i}{j}|{k}{m})" + assert np.isclose( + val, eri_gbasis[i, j, m, k], rtol=1e-10 + ), f"(ij|kl)!=(ij|lk) for ({i}{j}|{k}{m})" + assert np.isclose( + val, eri_gbasis[k, m, i, j], rtol=1e-10 + ), f"(ij|kl)!=(kl|ij) for ({i}{j}|{k}{m})" # --------------------------------------------------------------------------- # Tests: Improved (OS+HGP) implementation vs PySCF # --------------------------------------------------------------------------- -@pytest.mark.skipif(not HAS_PYSCF, reason="PySCF not installed") + class TestPySCFImproved: """Compare improved OS+HGP GBasis ERIs with PySCF reference.""" @@ -233,41 +255,61 @@ def test_h2_sto3g_chemist(self): basis, mol = make_h2_sto3g() eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") - np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, - err_msg="Improved H2/STO-3G chemist notation doesn't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Improved H2/STO-3G chemist notation doesn't match PySCF", + ) def test_h2_sto3g_physicist(self): """Test improved implementation H2/STO-3G in physicist notation.""" basis, mol = make_h2_sto3g() eri_gbasis = electron_repulsion_integral_improved(basis, notation="physicist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") eri_pyscf_physicist = eri_pyscf.transpose(0, 2, 1, 3) - np.testing.assert_allclose(eri_gbasis, eri_pyscf_physicist, rtol=1e-6, atol=1e-10, - err_msg="Improved H2/STO-3G physicist notation doesn't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf_physicist, + rtol=1e-6, + atol=1e-10, + err_msg="Improved H2/STO-3G physicist notation doesn't match PySCF", + ) def test_h2_primitive(self): """Test improved implementation with single primitive.""" basis, mol = make_h2_primitive() eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") - np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, - err_msg="Improved single primitive H2 ERIs don't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Improved single primitive H2 ERIs don't match PySCF", + ) def test_h2_sp_basis(self): """Test improved implementation H2 with s+p basis.""" basis, mol = make_h2_sp() eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") - np.testing.assert_allclose(eri_gbasis, eri_pyscf, rtol=1e-6, atol=1e-10, - err_msg="Improved H2 s+p basis ERIs don't match PySCF") + np.testing.assert_allclose( + eri_gbasis, + eri_pyscf, + rtol=1e-6, + atol=1e-10, + err_msg="Improved H2 s+p basis ERIs don't match PySCF", + ) def test_improved_matches_original(self): """Test that improved and original implementations agree on PySCF test case.""" @@ -276,8 +318,9 @@ def test_improved_matches_original(self): eri_old = electron_repulsion_integral(basis, notation="chemist") eri_new = electron_repulsion_integral_improved(basis, notation="chemist") - np.testing.assert_allclose(eri_new, eri_old, rtol=1e-10, - err_msg="Improved doesn't match original for H2/STO-3G") + np.testing.assert_allclose( + eri_new, eri_old, rtol=1e-10, err_msg="Improved doesn't match original for H2/STO-3G" + ) def test_improved_matches_original_sp(self): """Test that improved and original agree for s+p basis.""" @@ -286,15 +329,16 @@ def test_improved_matches_original_sp(self): eri_old = electron_repulsion_integral(basis, notation="chemist") eri_new = electron_repulsion_integral_improved(basis, notation="chemist") - np.testing.assert_allclose(eri_new, eri_old, rtol=1e-10, - err_msg="Improved doesn't match original for H2 s+p") + np.testing.assert_allclose( + eri_new, eri_old, rtol=1e-10, err_msg="Improved doesn't match original for H2 s+p" + ) # --------------------------------------------------------------------------- # Tests: Physical sanity checks # --------------------------------------------------------------------------- -@pytest.mark.skipif(not HAS_PYSCF, reason="PySCF not installed") + class TestPySCFPhysicalProperties: """Verify physical properties of ERIs using PySCF as cross-check.""" @@ -303,7 +347,7 @@ def test_positive_diagonal(self): basis, mol = make_h2_sto3g() eri_gbasis = electron_repulsion_integral_improved(basis, notation="chemist") - eri_pyscf = mol.intor('int2e') + eri_pyscf = mol.intor("int2e") n = eri_gbasis.shape[0] for i in range(n): @@ -316,15 +360,14 @@ def test_coulomb_greater_than_exchange(self): For same-atom shells: (00|00) >= (01|01) because the exchange integral involves orbital overlap which reduces the value. """ - basis, mol = make_h2_sto3g() + basis, _mol = make_h2_sto3g() eri = electron_repulsion_integral_improved(basis, notation="chemist") # Coulomb (00|00) should be > exchange (01|01) coulomb = eri[0, 0, 0, 0] exchange = eri[0, 1, 0, 1] - assert coulomb > exchange, \ - f"Coulomb ({coulomb:.6f}) should be > exchange ({exchange:.6f})" + assert coulomb > exchange, f"Coulomb ({coulomb:.6f}) should be > exchange ({exchange:.6f})" def test_eri_values_physically_reasonable(self): """Test that ERI values are in a physically reasonable range. @@ -341,13 +384,10 @@ def test_eri_values_physically_reasonable(self): # Diagonal integrals should be in reasonable range (0.1 - 2.0 a.u.) for i in range(eri.shape[0]): val = eri[i, i, i, i] - assert 0.01 < val < 10.0, \ - f"Diagonal integral ({i}{i}|{i}{i}) = {val} outside reasonable range" + assert ( + 0.01 < val < 10.0 + ), f"Diagonal integral ({i}{i}|{i}{i}) = {val} outside reasonable range" -if __name__ == '__main__': - if not HAS_PYSCF: - print("PySCF not installed. Install with: pip install pyscf") - print("Skipping PySCF comparison tests.") - else: - pytest.main([__file__, '-v']) +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 6e3a73eb9f87193ec24f8d9acc3f369f6400de44 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Wed, 11 Mar 2026 22:33:11 +0530 Subject: [PATCH 21/29] fix: use pytest.importorskip for pyscf instead of dev dependency --- pyproject.toml | 1 - tests/test_pyscf_comparison.py | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3dfa1618..82727343 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,6 @@ dev = [ "sphinx", "sphinx_autodoc_typehints", "sphinx-copybutton", - "pyscf>=1.6.1", ] doc = [ "numpydoc", diff --git a/tests/test_pyscf_comparison.py b/tests/test_pyscf_comparison.py index 78965159..fcded748 100644 --- a/tests/test_pyscf_comparison.py +++ b/tests/test_pyscf_comparison.py @@ -15,7 +15,6 @@ import numpy as np import pytest -from pyscf import gto from gbasis.contractions import GeneralizedContractionShell from gbasis.integrals.electron_repulsion import ( @@ -23,6 +22,8 @@ electron_repulsion_integral_improved, ) +gto = pytest.importorskip("pyscf.gto", reason="pyscf not installed") + # --------------------------------------------------------------------------- # Helper functions # --------------------------------------------------------------------------- From 0c215e64fc31af9b064f6b5c161ba9e1ece730ef Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Thu, 12 Mar 2026 20:13:16 +0530 Subject: [PATCH 22/29] fix: pass rho=harm_mean to boys_func to support erf/erfc attenuated potentials --- gbasis/integrals/_two_elec_int_improved.py | 2 +- gbasis/integrals/point_charge.py | 7 ++++--- tests/test_two_elec_int_improved.py | 4 ++-- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/gbasis/integrals/_two_elec_int_improved.py b/gbasis/integrals/_two_elec_int_improved.py index 627fd2a4..6ea46fa4 100644 --- a/gbasis/integrals/_two_elec_int_improved.py +++ b/gbasis/integrals/_two_elec_int_improved.py @@ -694,7 +694,7 @@ def compute_two_electron_integrals_os_hgp( 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 + integrals_m = boys_func(orders, weighted_dist[None, :, :, :, :], rho=harm_mean) * prefactor # --- Step 2: VRR --- integrals_vert = _vertical_recursion_relation( diff --git a/gbasis/integrals/point_charge.py b/gbasis/integrals/point_charge.py index 294f1bd2..b1597b46 100644 --- a/gbasis/integrals/point_charge.py +++ b/gbasis/integrals/point_charge.py @@ -1,9 +1,10 @@ """Module for computing point charge integrals.""" +import numpy as np +from scipy.special import hyp1f1 # pylint: disable=E0611 + from gbasis.base_two_symm import BaseTwoIndexSymmetric from gbasis.contractions import GeneralizedContractionShell from gbasis.integrals._one_elec_int import _compute_one_elec_integrals -import numpy as np -from scipy.special import hyp1f1 # pylint: disable=E0611 class PointChargeIntegral(BaseTwoIndexSymmetric): @@ -65,7 +66,7 @@ class PointChargeIntegral(BaseTwoIndexSymmetric): """ @staticmethod - def boys_func(orders, weighted_dist): + def boys_func(orders, weighted_dist, rho=None): r"""Return the value of Boys function for the given orders and weighted distances. The Coulombic Boys function can be written as a renormalized special case of the Kummer diff --git a/tests/test_two_elec_int_improved.py b/tests/test_two_elec_int_improved.py index 91f87ec2..5d9d9a1b 100644 --- a/tests/test_two_elec_int_improved.py +++ b/tests/test_two_elec_int_improved.py @@ -436,7 +436,7 @@ class TestPrimitiveScreening: def test_screening_zero_threshold(self): """With threshold=0, results match unscreened exactly.""" - def boys_func(orders, weighted_dist): + def boys_func(orders, weighted_dist, rho=None): return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) coord_a = np.array([0.0, 0.0, 0.0]) @@ -500,7 +500,7 @@ def boys_func(orders, weighted_dist): def test_screening_reasonable_threshold(self): """With reasonable threshold, results match within tolerance.""" - def boys_func(orders, weighted_dist): + def boys_func(orders, weighted_dist, rho=None): return hyp1f1(orders + 0.5, orders + 1.5, -weighted_dist) / (2 * orders + 1) coord_a = np.array([0.0, 0.0, 0.0]) From 874c15b660c332bdfc0d841a54ebbec7a419f0a8 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Thu, 12 Mar 2026 20:21:13 +0530 Subject: [PATCH 23/29] fix: replace tautological assertions with monotonicity check in test_threshold_sensitivity --- tests/test_schwarz_screening.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/test_schwarz_screening.py b/tests/test_schwarz_screening.py index 087a8004..333f1552 100644 --- a/tests/test_schwarz_screening.py +++ b/tests/test_schwarz_screening.py @@ -109,11 +109,9 @@ def test_threshold_sensitivity(self): # With loose threshold, should be significant sig_loose = shell_pair_significant(shell_a, shell_b, threshold=1e-30) - # The loose threshold should be at least as permissive - if not sig_strict: - assert sig_loose or not sig_loose # may or may not be significant - if sig_loose: - assert True # loose is at least as permissive + # Monotonicity: any pair significant under strict threshold must also be + # significant under loose threshold (smaller threshold = more permissive) + assert (not sig_strict) or sig_loose class TestSchwarzBounds: From 65879ade82218ce91d77103160fad65d44f37082 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Thu, 12 Mar 2026 20:28:25 +0530 Subject: [PATCH 24/29] fix: pass screener via kwargs instead of class-level state for thread-safety --- gbasis/integrals/electron_repulsion.py | 58 +++++++++++++------------- 1 file changed, 30 insertions(+), 28 deletions(-) diff --git a/gbasis/integrals/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index a48b5910..0ab5a3cd 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -276,11 +276,9 @@ class ElectronRepulsionIntegralImproved(BaseFourIndexSymmetric): """ boys_func = PointChargeIntegral.boys_func - _screener = None - _contraction_index_map = None @classmethod - def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four): + def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four, **kwargs): r"""Return electron-electron repulsion integral using the OS+HGP algorithm. Parameters @@ -316,13 +314,15 @@ def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four): raise TypeError("`cont_four` must be a `GeneralizedContractionShell` instance.") # --- Schwarz screening: skip negligible shell quartets --- - if cls._screener is not None and cls._contraction_index_map is not None: - i = cls._contraction_index_map.get(id(cont_one)) - j = cls._contraction_index_map.get(id(cont_two)) - k = cls._contraction_index_map.get(id(cont_three)) - l_idx = cls._contraction_index_map.get(id(cont_four)) + screener = kwargs.get("screener") + index_map = kwargs.get("contraction_index_map") + if screener is not None and index_map is not None: + i = index_map.get(id(cont_one)) + j = index_map.get(id(cont_two)) + k = index_map.get(id(cont_three)) + l_idx = index_map.get(id(cont_four)) if i is not None and j is not None and k is not None and l_idx is not None: - if not cls._screener.is_significant(i, j, k, l_idx): + if not screener.is_significant(i, j, k, l_idx): shape = ( cont_one.coeffs.shape[1], len(cont_one.angmom_components_cart), @@ -441,6 +441,7 @@ def electron_repulsion_integral_improved( if notation not in ["physicist", "chemist"]: raise ValueError("`notation` must be one of 'physicist' or 'chemist'") + screening_kwargs = {} if schwarz_threshold > 0: index_map = {id(shell): i for i, shell in enumerate(basis)} screener = SchwarzScreener( @@ -449,25 +450,26 @@ def electron_repulsion_integral_improved( compute_two_electron_integrals_os_hgp, schwarz_threshold, ) - ElectronRepulsionIntegralImproved._screener = screener - ElectronRepulsionIntegralImproved._contraction_index_map = index_map - - try: - 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) - finally: - ElectronRepulsionIntegralImproved._screener = None - ElectronRepulsionIntegralImproved._contraction_index_map = None + screening_kwargs = {"screener": screener, "contraction_index_map": index_map} + + coord_type = [ct for ct in [shell.coord_type for shell in basis]] + + if transform is not None: + array = ElectronRepulsionIntegralImproved(basis).construct_array_lincomb( + transform, coord_type, **screening_kwargs + ) + elif all(ct == "cartesian" for ct in coord_type): + array = ElectronRepulsionIntegralImproved(basis).construct_array_cartesian( + **screening_kwargs + ) + elif all(ct == "spherical" for ct in coord_type): + array = ElectronRepulsionIntegralImproved(basis).construct_array_spherical( + **screening_kwargs + ) + else: + array = ElectronRepulsionIntegralImproved(basis).construct_array_mix( + coord_type, **screening_kwargs + ) if notation == "physicist": array = np.transpose(array, (0, 2, 1, 3)) From fb9e1a6bd1d266e47526f7cdeed4d3c950e57cce Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Thu, 12 Mar 2026 20:36:46 +0530 Subject: [PATCH 25/29] perf: cache SchwarzScreener to avoid recomputing bounds on every call --- gbasis/integrals/electron_repulsion.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/gbasis/integrals/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index 0ab5a3cd..dc232f21 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -8,6 +8,9 @@ from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp from gbasis.integrals.point_charge import PointChargeIntegral +# Cache for Schwarz screeners keyed by (basis_ids, threshold) +_schwarz_screener_cache = {} + class ElectronRepulsionIntegral(BaseFourIndexSymmetric): """Class for constructing electron-electron repulsion integrals. @@ -443,13 +446,16 @@ def electron_repulsion_integral_improved( screening_kwargs = {} if schwarz_threshold > 0: + cache_key = (tuple(id(shell) for shell in basis), schwarz_threshold) + if cache_key not in _schwarz_screener_cache: + _schwarz_screener_cache[cache_key] = SchwarzScreener( + list(basis), + ElectronRepulsionIntegralImproved.boys_func, + compute_two_electron_integrals_os_hgp, + schwarz_threshold, + ) + screener = _schwarz_screener_cache[cache_key] index_map = {id(shell): i for i, shell in enumerate(basis)} - screener = SchwarzScreener( - list(basis), - ElectronRepulsionIntegralImproved.boys_func, - compute_two_electron_integrals_os_hgp, - schwarz_threshold, - ) screening_kwargs = {"screener": screener, "contraction_index_map": index_map} coord_type = [ct for ct in [shell.coord_type for shell in basis]] From bfac8b09e434b79af97e35fe84ebdef82b393ac1 Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Fri, 13 Mar 2026 22:58:02 +0530 Subject: [PATCH 26/29] revert: restore original ElectronRepulsionIntegral to use old algorithm --- gbasis/integrals/electron_repulsion.py | 67 +++++++++++++++++--------- 1 file changed, 44 insertions(+), 23 deletions(-) diff --git a/gbasis/integrals/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index dc232f21..e9194158 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -5,6 +5,10 @@ from gbasis.base_four_symm import BaseFourIndexSymmetric from gbasis.contractions import GeneralizedContractionShell from gbasis.integrals._schwarz_screening import SchwarzScreener +from gbasis.integrals._two_elec_int import ( + _compute_two_elec_integrals, + _compute_two_elec_integrals_angmom_zero, +) from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp from gbasis.integrals.point_charge import PointChargeIntegral @@ -159,29 +163,46 @@ 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 - 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, - ) + 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 = np.transpose(integrals, (4, 0, 5, 1, 6, 2, 7, 3)) # TODO: if we swap the contractions, we need to unswap them here From f2f0e0d137125ea2ba7fdceb8d315849df1a6d5a Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Sat, 14 Mar 2026 01:23:07 +0530 Subject: [PATCH 27/29] fix: use get_boys_function from boys_functions module in ElectronRepulsionIntegralImproved --- gbasis/integrals/electron_repulsion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gbasis/integrals/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index e9194158..37ad05b5 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -10,6 +10,7 @@ _compute_two_elec_integrals_angmom_zero, ) from gbasis.integrals._two_elec_int_improved import compute_two_electron_integrals_os_hgp +from gbasis.integrals.boys_functions import get_boys_function from gbasis.integrals.point_charge import PointChargeIntegral # Cache for Schwarz screeners keyed by (basis_ids, threshold) @@ -299,7 +300,7 @@ class ElectronRepulsionIntegralImproved(BaseFourIndexSymmetric): """ - boys_func = PointChargeIntegral.boys_func + boys_func = get_boys_function("coulomb") @classmethod def construct_array_contraction(cls, cont_one, cont_two, cont_three, cont_four, **kwargs): From 76107eaa50b0834f71965b2829c81ce62a6f1adf Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Sat, 14 Mar 2026 01:40:12 +0530 Subject: [PATCH 28/29] fix: use get_boys_function directly instead of ElectronRepulsionIntegralImproved.boys_func --- gbasis/integrals/electron_repulsion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gbasis/integrals/electron_repulsion.py b/gbasis/integrals/electron_repulsion.py index 37ad05b5..295b0cfa 100644 --- a/gbasis/integrals/electron_repulsion.py +++ b/gbasis/integrals/electron_repulsion.py @@ -472,7 +472,7 @@ def electron_repulsion_integral_improved( if cache_key not in _schwarz_screener_cache: _schwarz_screener_cache[cache_key] = SchwarzScreener( list(basis), - ElectronRepulsionIntegralImproved.boys_func, + get_boys_function("coulomb"), compute_two_electron_integrals_os_hgp, schwarz_threshold, ) From 104446c332230df323e66ebda2b0ed4ec2a0e5fc Mon Sep 17 00:00:00 2001 From: Aayush Gupta Date: Sat, 14 Mar 2026 01:58:21 +0530 Subject: [PATCH 29/29] fix: use boys_functions module directly instead of inline hyp1f1 or class reference --- gbasis/integrals/point_charge.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/gbasis/integrals/point_charge.py b/gbasis/integrals/point_charge.py index b1597b46..fe2a7cd9 100644 --- a/gbasis/integrals/point_charge.py +++ b/gbasis/integrals/point_charge.py @@ -1,10 +1,11 @@ """Module for computing point charge integrals.""" + import numpy as np -from scipy.special import hyp1f1 # pylint: disable=E0611 from gbasis.base_two_symm import BaseTwoIndexSymmetric from gbasis.contractions import GeneralizedContractionShell from gbasis.integrals._one_elec_int import _compute_one_elec_integrals +from gbasis.integrals.boys_functions import boys_function_standard class PointChargeIntegral(BaseTwoIndexSymmetric): @@ -114,7 +115,7 @@ def boys_func(orders, weighted_dist, rho=None): side. """ - return hyp1f1(orders + 1 / 2, orders + 3 / 2, -weighted_dist) / (2 * orders + 1) + return boys_function_standard(orders, weighted_dist) @classmethod def construct_array_contraction(