Speed up the test suite without dropping coverage

[aarmey]

Summary

• Correctness fix: stateCommon.py grabbed the wrong Cython fused-function overload for the digamma function (psi) — __pyx_fuse_0psi is psi's complex128 specialization, not the real-valued one (__pyx_fuse_1psi). Calling the complex overload through a real-valued ctypes.CFUNCTYPE reads garbage for the argument's missing imaginary half. This silently biased gamma_mle_closed_form's shape-parameter MLE (its "exact, instantaneous" uncensored fast path — e.g. a synthetic Gamma(4, 2) sample recovered shape 3.84 instead of 3.98) and threw off gamma_estimator's analytic gradient by ~5-10x in the shape direction for censored fits. That bad gradient is what was causing SLSQP to hit its 200-iteration cap on otherwise well-posed problems and fall back to the much slower trust-constr solver, which was the dominant remaining cost in test_cv and test_works_through_crossval. Root-caused by comparing the JIT-compiled function's output against its own .py_func (pure-Python) fallback, which disagreed — see commit for the full trace.
• Rewrote StateDistribution.division_probability()'s integrand to use scipy.special (gammaln/gammaincc) directly instead of scipy.stats frozen-distribution pdf/sf, which was ~100x more expensive per call inside the quad() integration used every M-step.
• Rewrote StateDistribution.logpdf() the same way — it runs once per state on every E step over the whole lineage array, and was still going through the frozen rv_continuous machinery. Verified the new output matches the old to float precision. Cuts test_works_through_crossval from ~14.7s to ~6.3s in isolation on top of the division_probability fix.
• Enabled pytest-xdist (-n auto, already a dev dependency but unused) via addopts in pyproject.toml, and pinned OMP_NUM_THREADS/OPENBLAS_NUM_THREADS/MKL_NUM_THREADS/NUMEXPR_NUM_THREADS to 1 in a new lineage/tests/conftest.py so xdist workers don't fight scipy's own internal BLAS threading for cores.
• Combined, the full suite drops from ~125s to ~16s wall clock, with no tests removed, skipped, or weakened — and the shape-parameter MLE is now unbiased.

Remaining known costs (not addressed here)

• Numba JIT compilation of gamma_estimator's njit functions pays a one-time ~3-4s cost per process the first time each type signature is hit; under -n auto several workers re-pay it independently. Would need cache=True on the njit decorators to amortize across processes/runs.
• The Baum-Welch E-step recursion (HMM/E_step.py:get_beta_and_NF) is plain Python/numpy, not JIT-accelerated, and is now one of the larger remaining per-test costs (~7s of test_cv[0]'s ~12.5s in isolation). This is legitimate algorithmic cost rather than a bug, and a bigger undertaking to address than the leaf-function fixes above.

Test plan

• uv run pytest — all tests pass (67, including master's new Hypothesis tests) in ~16s, down from ~125s on the original branch point
• uv run ruff check on touched files — passes
• Verified the psi fix numerically: compared the JIT function against its .py_func pure-Python twin on identical inputs (they now agree; they didn't before), and checked __pyx_fuse_1psi against scipy.special.psi directly across several values
• Verified logpdf()'s new output matches the old frozen-distribution implementation to float precision (np.testing.assert_allclose(..., rtol=1e-10, atol=1e-10)) on synthetic data including censored/unknown-fate rows
• Spot-checked division_probability() still matches the existing test_sub_densities_sum_to_one / test_estimator_recovers_parameters / test_phase_estimator_recovers_parameters correctness checks

🤖 Generated with Claude Code