fix: handle different K per center in _optimized_contraction

Author: San1357

gbasis/integrals/_two_elec_int_improved.py

@@ -15,79 +15,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]
+    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_a,
-    exps_b,
-    exps_c,
-    exps_d,
-    coeffs_a,
-    coeffs_b,
-    coeffs_c,
-    coeffs_d,
-    angmom_a,
-    angmom_b,
-    angmom_c,
-    angmom_d,
-):
+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 : 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).
     """
-    # 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]
+    # Compute norms per center (supports different K per center)
+    norms = [((2 / np.pi) * e) ** 0.75 * (4 * e) ** (ang / 2) for e, ang in zip(exps, angmoms)]
+    coeffs_norm = [c * n[:, np.newaxis] for c, n in zip(coeffs, norms)]
+    coeffs_a_norm, coeffs_b_norm, coeffs_c_norm, coeffs_d_norm = coeffs_norm
 
     # Use einsum with optimization for contraction
     # Input: (c_x, c_y, c_z, a_x, a_y, a_z, K_d, K_b, K_c, K_a)
@@ -594,10 +570,18 @@ def _horizontal_recursion_relation(
     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)
+    norm_a = _get_factorial2_norm(tuple(map(tuple, angmom_components_a))).reshape(
+        -1, 1, 1, 1, 1, 1, 1, 1
+    )
+    norm_b = _get_factorial2_norm(tuple(map(tuple, angmom_components_b))).reshape(
+        1, -1, 1, 1, 1, 1, 1, 1
+    )
+    norm_c = _get_factorial2_norm(tuple(map(tuple, angmom_components_c))).reshape(
+        1, 1, -1, 1, 1, 1, 1, 1
+    )
+    norm_d = _get_factorial2_norm(tuple(map(tuple, angmom_components_d))).reshape(
+        1, 1, 1, -1, 1, 1, 1, 1
+    )
 
     integrals = integrals * norm_a * norm_b * norm_c * norm_d
 
@@ -725,18 +709,9 @@ def compute_two_electron_integrals_os_hgp(
     # --- 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,
+        (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) ---

tests/test_two_elec_int_improved.py

@@ -13,6 +13,7 @@
 from gbasis.integrals._two_elec_int_improved import (
     _electron_transfer_recursion,
     _get_factorial2_norm,
+    _optimized_contraction,
     _vertical_recursion_relation,
     compute_two_electron_integrals_os_hgp,
 )
@@ -161,29 +162,59 @@ 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."""
 
     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)