diff --git a/lineage/states/StateDistributionCR.py b/lineage/states/StateDistributionCR.py index 19e340a5e..267a85b89 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 @@ -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 diff --git a/lineage/states/stateCommon.py b/lineage/states/stateCommon.py index f34137fcf..84cedf7f7 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_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) 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 7d0a6b84c..6b85b0395 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,5 +73,6 @@ minversion = "6.0" testpaths = [ "lineage/tests", ] +addopts = "-n auto"