Speed up the test suite without dropping coverage (#1011)

* Speed up the test suite without dropping coverage

The two slowest tests spent most of their time in
StateDistribution.division_probability(), which numerically integrates a
product of scipy.stats frozen-distribution pdf/sf calls; that generic
rv_continuous machinery costs far more than the arithmetic itself when
called ~200 times per quad() evaluation. Rewriting the integrand against
scipy.special (gammaln/gammaincc) directly cuts
test_works_through_crossval from ~45s to ~18s with identical results.

Also enable pytest-xdist (-n auto, already a dev dependency but unused)
and pin BLAS/OpenMP thread counts to 1 in a new tests/conftest.py so
xdist's worker processes don't fight scipy's own internal threading for
cores. Combined, the full suite drops from ~125s to ~47s.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Fix wrong Cython fused-function overload for digamma (psi)

lineage/states/stateCommon.py grabbed the raw address of
scipy.special.cython_special.__pyx_fuse_0psi and called it through a
ctypes CFUNCTYPE(c_double, c_double). That symbol is psi's *complex128*
specialization ("__pyx_t_double_complex (__pyx_t_double_complex, int)");
the real-valued one used here is __pyx_fuse_1psi. Calling the complex
overload through a real signature reads garbage for the missing
imaginary half of the argument. It happened to look right under plain
ctypes calls (the unused slot came back zero by luck), which is why this
went unnoticed, but numba's own generated call sequence for the same
function pointer leaves that slot non-zero, so every njit call to psi()
silently returned a wrong value.

This directly biased two things:
- gamma_mle_closed_form's Newton-Raphson solve for the shape parameter
  (its "100% exact, instantaneous" uncensored fast path), e.g. a
  synthetic Gamma(4, 2) sample recovered shape 3.84 instead of 3.98.
- gamma_estimator's analytic gradient for uncensored observations,
  which was off by roughly 5-10x in the shape-parameter direction. That
  bad gradient is why SLSQP kept hitting its 200-iteration cap on
  otherwise well-posed censored fits and falling back to the much
  slower trust-constr solver, which was the dominant cost in
  test_cv[3] and test_works_through_crossval.

Fixing the symbol lookup resolves both: the full test suite drops from
~47s to ~22s on top of the earlier speedups, and the shape MLE is now
unbiased.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Vectorize StateDistributionCR.logpdf against scipy.special

logpdf() runs once per state on every E step over the whole lineage
array, but still went through the frozen rv_continuous machinery
(self.div_clock.logpdf/logsf, self.death_clock.logpdf/logsf) for that -
the same class of overhead already fixed in division_probability(),
just not applied here. Replaced with direct vectorized formulas against
scipy.special (gammaln/gammaincc), matching the old output to float
precision (verified numerically) while cutting test_works_through_crossval
from ~14.7s to ~6.3s in isolation.

The div_clock/death_clock properties are left in place since
compare_emissions.py still uses them outside this hot path.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: aarmey

lineage/states/StateDistributionCR.py

@@ -30,6 +30,7 @@
 import scipy.stats as sp
 from scipy.integrate import quad
 from scipy.sparse import csr_array
+from scipy.special import gammaincc, gammaln
 
 from .stateCommon import censor_lineage_gamma, censor_lineage_gaphs, gamma_estimator
 
@@ -38,6 +39,16 @@
 TIME_FLOOR = 1e-10
 
 
+def _gamma_logpdf(t: np.ndarray, a: float, scale: float) -> np.ndarray:
+    """log Gamma(a, scale) pdf, vectorized directly against scipy.special."""
+    return (a - 1.0) * np.log(t) - t / scale - a * np.log(scale) - gammaln(a)
+
+
+def _gamma_logsf(t: np.ndarray, a: float, scale: float) -> np.ndarray:
+    """log Gamma(a, scale) survival function, vectorized directly against scipy.special."""
+    return np.log(gammaincc(a, t / scale))
+
+
 def event_masks(x: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Split cells into the three competing-risks cases.
 
@@ -197,13 +208,23 @@ def death_clock(self):
         return sp.gamma(a=self.params[3], scale=self.params[4])
 
     def division_probability(self) -> float:
-        """P(the division clock fires first) = int f_D(t) S_X(t) dt."""
-        div, death = self.div_clock, self.death_clock
+        """P(the division clock fires first) = int f_D(t) S_X(t) dt.
+
+        ``quad`` evaluates the integrand at ~100-200 points per call, and this runs
+        once per state per M step. Using the frozen ``rv_continuous`` machinery for
+        each point (as ``div.pdf`` / ``death.sf`` do) costs orders of magnitude more
+        than the arithmetic itself, so the integrand is written directly against
+        ``scipy.special`` instead.
+        """
+        a1, s1 = self.params[1], self.params[2]
+        a2, s2 = self.params[3], self.params[4]
+        log_norm = -gammaln(a1) - a1 * np.log(s1)
 
         def integrand(t):
-            return div.pdf(t) * death.sf(t)
+            logpdf = log_norm + (a1 - 1.0) * np.log(t) - t / s1
+            return np.exp(logpdf) * gammaincc(a2, t / s2)
 
-        upper = div.ppf(1.0 - 1e-9)
+        upper = self.div_clock.ppf(1.0 - 1e-9)
         val = quad(integrand, 0.0, upper, limit=100)[0]
         return float(np.clip(val, 0.0, 1.0))
 
@@ -231,19 +252,26 @@ def dof(self) -> int:
         return 3 if self.fixed_death_shape else 4
 
     def logpdf(self, x: np.ndarray) -> np.ndarray:
-        """Competing-risks log likelihood of each cell's phase observation."""
+        """Competing-risks log likelihood of each cell's phase observation.
+
+        Called once per state on every E step, over the whole lineage array, so it is
+        vectorized directly against ``scipy.special`` rather than going through the
+        frozen ``rv_continuous`` machinery (``self.div_clock.logpdf`` etc.), which pays
+        generic validation/broadcasting overhead on every call regardless of array size.
+        """
         divided, died, censored = event_masks(x)
         timed = divided | died | censored
         t = np.clip(x[:, 1], TIME_FLOOR, None)
 
-        div, death = self.div_clock, self.death_clock
+        a1, s1 = self.params[1], self.params[2]
+        a2, s2 = self.params[3], self.params[4]
 
         ll = np.zeros(x.shape[0])
         # Every timed cell survived both clocks up to t; the one that fired then
         # swaps its survival term for a density.
-        ll[timed] += div.logsf(t[timed]) + death.logsf(t[timed])
-        ll[divided] += div.logpdf(t[divided]) - div.logsf(t[divided])
-        ll[died] += death.logpdf(t[died]) - death.logsf(t[died])
+        ll[timed] += _gamma_logsf(t[timed], a1, s1) + _gamma_logsf(t[timed], a2, s2)
+        ll[divided] += _gamma_logpdf(t[divided], a1, s1) - _gamma_logsf(t[divided], a1, s1)
+        ll[died] += _gamma_logpdf(t[died], a2, s2) - _gamma_logsf(t[died], a2, s2)
 
         assert not np.any(np.isnan(ll))
         return ll

lineage/states/stateCommon.py

@@ -182,7 +182,11 @@ def bern_estimator(bern_obs: np.ndarray, gammas: np.ndarray):
 _addr_gammaln = get_cython_function_address("scipy.special.cython_special", "gammaln")
 gammaln = CFUNCTYPE(c_double, c_double)(_addr_gammaln)
 
-_addr_psi = get_cython_function_address("scipy.special.cython_special", "__pyx_fuse_0psi")
+# psi is a fused Cython function with both a complex and a real specialization;
+# __pyx_fuse_0psi is the complex128 one (silently reads garbage for the missing
+# imaginary half when called through this real-valued CFUNCTYPE) and __pyx_fuse_1psi
+# is the real double one we actually want.
+_addr_psi = get_cython_function_address("scipy.special.cython_special", "__pyx_fuse_1psi")
 psi = CFUNCTYPE(c_double, c_double)(_addr_psi)
 
 

lineage/tests/conftest.py

@@ -0,0 +1,13 @@
+"""Test-session setup.
+
+Several tests drive scipy/numpy linear algebra hard enough that BLAS's own thread
+pool competes with pytest-xdist's worker processes for cores, which slows the suite
+down rather than speeding it up. Pin each library to a single BLAS thread so the
+parallelism comes from xdist alone; this must happen before numpy/scipy are
+imported anywhere, so it lives at the top of the first conftest.py loaded.
+"""
+
+import os
+
+for _var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS"):
+    os.environ.setdefault(_var, "1")

pyproject.toml

@@ -73,5 +73,6 @@ minversion = "6.0"
 testpaths = [
     "lineage/tests",
 ]
+addopts = "-n auto"