From ae120364ba4ec5241e8c18f33a07085022798ad7 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 08:06:19 -0700 Subject: [PATCH 1/3] 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 --- lineage/states/StateDistributionCR.py | 19 +++++++++++++++---- lineage/tests/conftest.py | 13 +++++++++++++ pyproject.toml | 1 + 3 files changed, 29 insertions(+), 4 deletions(-) create mode 100644 lineage/tests/conftest.py diff --git a/lineage/states/StateDistributionCR.py b/lineage/states/StateDistributionCR.py index 19e340a5e..cc45d4e87 100644 --- a/lineage/states/StateDistributionCR.py +++ b/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 @@ -197,13 +198,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)) diff --git a/lineage/tests/conftest.py b/lineage/tests/conftest.py new file mode 100644 index 000000000..f1799421c --- /dev/null +++ b/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") diff --git a/pyproject.toml b/pyproject.toml index aefb5984a..5b380f209 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,5 +67,6 @@ minversion = "6.0" testpaths = [ "lineage/tests", ] +addopts = "-n auto" From 2671208526f5f3b64a654f5662a95cc915a19df6 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 08:16:55 -0700 Subject: [PATCH 2/3] 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 --- lineage/states/stateCommon.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/lineage/states/stateCommon.py b/lineage/states/stateCommon.py index b518f4980..45d58f321 100644 --- a/lineage/states/stateCommon.py +++ b/lineage/states/stateCommon.py @@ -182,7 +182,11 @@ def bern_estimator(bern_obs: np.ndarray, gammas: np.ndarray): addr = get_cython_function_address("scipy.special.cython_special", "gammaln") gammaln = CFUNCTYPE(c_double, c_double)(addr) -addr = 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 = get_cython_function_address("scipy.special.cython_special", "__pyx_fuse_1psi") psi = CFUNCTYPE(c_double, c_double)(addr) From cc75ed2e5e91160e87c9e39ec2938700a3dab997 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 08:25:13 -0700 Subject: [PATCH 3/3] 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 --- lineage/states/StateDistributionCR.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/lineage/states/StateDistributionCR.py b/lineage/states/StateDistributionCR.py index cc45d4e87..267a85b89 100644 --- a/lineage/states/StateDistributionCR.py +++ b/lineage/states/StateDistributionCR.py @@ -39,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. @@ -242,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