refactor to distinguish between jax & scipy double fact

Author: lazyoracle

src/isotropic/e2.py

@@ -6,28 +6,9 @@
 import jax.random as random
 from jax import Array
 from jax.typing import ArrayLike
-from scipy.special import factorial2
 
 from isotropic.utils.bisection import get_theta
-
-
-def double_factorial_ratio(n: int, m: int) -> float:
-    """
-    Compute the ratio of double factorials n!! / m!!.
-
-    Parameters
-    ----------
-    n : int
-        The numerator double factorial.
-    m : int
-        The denominator double factorial.
-
-    Returns
-    -------
-    float
-        The ratio of the double factorials.
-    """
-    return factorial2(n) / factorial2(m)
+from isotropic.utils.distribution import double_factorial_ratio_scipy
 
 
 def F_j(theta_j: float, j: int, d: int) -> Array:
@@ -49,10 +30,10 @@ def F_j(theta_j: float, j: int, d: int) -> Array:
         The value of the function F_j evaluated at theta_j.
     """
     dj = d - j
-    numoverden = double_factorial_ratio(dj - 2, dj - 1)
+    numoverden = double_factorial_ratio_scipy(dj - 2, dj - 1)
 
     def F_odd(_):
-        C_j = (1 / 2) * double_factorial_ratio(dj - 1, dj - 2)
+        C_j = (1 / 2) * double_factorial_ratio_scipy(dj - 1, dj - 2)
         prefactor = C_j * numoverden
         k_max = (dj - 2) // 2  # upper bound for k in range
         k_vals = jnp.arange(0, k_max + 1)
@@ -72,7 +53,7 @@ def product_term(k):
         return prefactor - C_j * jnp.cos(theta_j) * sum_terms
 
     def F_even(_):
-        C_j = (1 / jnp.pi) * double_factorial_ratio(dj - 1, dj - 2)
+        C_j = (1 / jnp.pi) * double_factorial_ratio_scipy(dj - 1, dj - 2)
         prefactor = C_j * numoverden * theta_j
         k_max = (dj - 1) // 2
         k_vals = jnp.arange(1, k_max + 1)

src/isotropic/utils/distribution.py

@@ -2,9 +2,29 @@
 
 import jax.numpy as jnp
 from jax import Array
+from scipy.special import factorial2
 
 
-def double_factorial(n: int) -> Array:
+def double_factorial_ratio_scipy(num: int, den: int) -> float:
+    """
+    Compute the ratio of double factorials num!! / den!!.
+
+    Parameters
+    ----------
+    num : int
+        The numerator double factorial.
+    den : int
+        The denominator double factorial.
+
+    Returns
+    -------
+    float
+        The ratio of the double factorials.
+    """
+    return factorial2(num) / factorial2(den)
+
+
+def double_factorial_jax(n: int) -> Array:
     """
     Helper function to compute double factorial:
 
@@ -24,7 +44,7 @@ def double_factorial(n: int) -> Array:
     return jnp.where(n <= 0, 1, jnp.prod(jnp.arange(n, 0, -2, dtype=jnp.uint64)))
 
 
-def double_factorial_ratio(num: int, den: int) -> Array:
+def double_factorial_ratio_jax(num: int, den: int) -> Array:
     """
     Computes the ratio of double factorials:
 
@@ -101,8 +121,8 @@ def normal_integrand(theta: float, d: int, sigma: float) -> Array:
     # sigma = jnp.asarray(sigma)
 
     # factorial components
-    numerator_factorial = double_factorial(d - 1)
-    denominator_factorial = double_factorial(d - 2)
+    numerator_factorial = factorial2(d - 1)
+    denominator_factorial = factorial2(d - 2)
 
     # Numerator components
     one_minus_sigma_sq = 1.0 - sigma**2

tests/test_utils.py

@@ -4,8 +4,8 @@
 
 from isotropic.utils.bisection import get_theta
 from isotropic.utils.distribution import (
-    double_factorial,
-    double_factorial_ratio,
+    double_factorial_jax,
+    double_factorial_ratio_jax,
     normal_integrand,
 )
 from isotropic.utils.linalg import jax_null_space
@@ -58,42 +58,42 @@ def F(theta):
     )
 
 
-def test_double_factorial():
+def test_double_factorial_jax():
     # Test even double factorial
     n_even = 6
-    result_even = double_factorial(n_even)
+    result_even = double_factorial_jax(n_even)
     expected_even = factorial2(n_even)
     assert jnp.isclose(result_even, expected_even), (
         f"Expected {expected_even}, got {result_even}"
     )
 
     # Test odd double factorial
     n_odd = 5
-    result_odd = double_factorial(n_odd)
+    result_odd = double_factorial_jax(n_odd)
     expected_odd = factorial2(n_odd)
     assert jnp.isclose(result_odd, expected_odd), (
         f"Expected {expected_odd}, got {result_odd}"
     )
 
     # Test zero double factorial
     n_zero = 0
-    result_zero = double_factorial(n_zero)
+    result_zero = double_factorial_jax(n_zero)
     expected_zero = factorial2(n_zero)
     assert jnp.isclose(result_zero, expected_zero), (
         f"Expected {expected_zero}, got {result_zero}"
     )
 
 
-def test_double_factorial_ratio():
+def test_double_factorial_ratio_jax():
     num, den = (2**8) - 1, (2**8) - 2
-    ratio_received = double_factorial_ratio(num, den)
+    ratio_received = double_factorial_ratio_jax(num, den)
     ratio_expected = factorial2(num) / factorial2(den)
     assert jnp.isclose(ratio_received, ratio_expected), (
         f"Expected {ratio_expected}, got {ratio_received}"
     )
 
     num, den = (2**8) - 3, (2**8) - 1
-    ratio_received = double_factorial_ratio(num, den)
+    ratio_received = double_factorial_ratio_jax(num, den)
     ratio_expected = factorial2(num) / factorial2(den)
     assert jnp.isclose(ratio_received, ratio_expected), (
         f"Expected {ratio_expected}, got {ratio_received}"