From dabae14271f60c6913d7f4f6f4f8348d5f6e6ea5 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 6 Sep 2026 20:13:17 -0700 Subject: [PATCH 1/6] Add competing-risks emission for cell lifetime and fate The Gamma/Bernoulli emissions treat a cell's fate and its phase duration as independent, and zero out the duration of any cell that dies, so every death time in the dataset is discarded. They also score a time-censored cell with P(division > t) when the correct statement is P(division > t AND death > t), which leaves the emission unnormalized whenever there is censoring. StateDistributionCR models each phase as a division clock racing a death clock and observes the first to fire. Death times now carry information, the censored term is correct, and the division probability is derived from the two clocks rather than fit as a free Bernoulli parameter. The G1 death clock is pinned to a constant hazard and G2's shape is free, which matches the data: pooled G1 death times give a Weibull shape of 1.03 while G2 gives 2.75, consistently across every drug and concentration. The flat params array keeps the Gamma/GaPhs layout in its leading entries so existing figure code that indexes it positionally is unaffected. --- lineage/BaumWelch.py | 5 +- lineage/compare_emissions.py | 235 +++++++++++++++ lineage/states/StateDistributionCR.py | 336 ++++++++++++++++++++++ lineage/tests/test_StateDistributionCR.py | 160 +++++++++++ 4 files changed, 735 insertions(+), 1 deletion(-) create mode 100644 lineage/compare_emissions.py create mode 100644 lineage/states/StateDistributionCR.py create mode 100644 lineage/tests/test_StateDistributionCR.py diff --git a/lineage/BaumWelch.py b/lineage/BaumWelch.py index 0f437547b..1e064b023 100644 --- a/lineage/BaumWelch.py +++ b/lineage/BaumWelch.py @@ -7,7 +7,7 @@ from .HMM.E_step import get_beta_and_NF, get_gamma, get_MSD from .HMM.M_step import get_all_zetas, sum_nonleaf_gammas from .LineageTree import get_Emission_Likelihoods -from .states.StateDistributionGamma import atonce_estimator +from .states.StateDistributionGamma import atonce_estimator as gamma_atonce_estimator from .tHMM import tHMM @@ -220,6 +220,9 @@ def do_M_E_step_atonce(all_tHMMobj: list[tHMM], all_gammas: list[list[np.ndarray else: cells.append(all_cells) + # Emission classes may supply their own at-once estimator; fall back to the Gamma one. + atonce_estimator = getattr(all_tHMMobj[0].estimate.E[0], "atonce_estimator", gamma_atonce_estimator) + # reshape the gammas so that each list in this list of lists is for each state. if phase: atonce_estimator(all_tHMMobj, G1cells, gms, "G1") # [shape, scale1, scale2, scale3, scale4] for G1 diff --git a/lineage/compare_emissions.py b/lineage/compare_emissions.py new file mode 100644 index 000000000..18111b6f1 --- /dev/null +++ b/lineage/compare_emissions.py @@ -0,0 +1,235 @@ +"""Cross-validated comparison of the Gamma/Bernoulli and competing-risks emissions. + +The two emissions are densities over *different* observation spaces. The competing-risks +form in :mod:`lineage.states.StateDistributionCR` puts a density on when a cell died, +which the Gamma/Bernoulli form has no way to express -- it scores a death as a bare +Bernoulli outcome and discards the time. Comparing their raw likelihoods would charge +the competing-risks model for predicting strictly more. + +The headline metric therefore coarsens death timing away, leaving a space both models +describe: + + ============ ========================= ==================== + outcome competing risks Gamma/Bernoulli + ============ ========================= ==================== + transition t ``f_D(t) S_X(t)`` ``p f_D(t)`` + death ``1 - P(divide)`` ``1 - p`` + censored, c ``S_D(c) S_X(c)`` ``S_D(c)`` + ============ ========================= ==================== + +The transition branches carry identical total mass, so that comparison is like for +like. The censored branch is where the Gamma/Bernoulli form ignores that a censored +cell also did not die, and so claims more probability than it is entitled to -- see +:func:`outcome_mass`, which shows its total exceeding one by up to ~10% for the +short-horizon cells that make up much of this data. That bias runs in the Gamma +model's favour, which makes the comparison conservative. + +Run as ``python -m lineage.compare_emissions ``, e.g. +``python -m lineage.compare_emissions AllLapatinib 2,3,4 5``. +""" + +import json +import sys +import time + +import numpy as np +import scipy.stats as sp +from scipy.special import logsumexp + +from .Analyze import Analyze_list +from .LineageTree import LineageTree +from .states.StateDistributionCR import StateDistributionPhase as CR +from .states.StateDistributionCR import event_masks +from .states.StateDistributionGaPhs import StateDistribution as GA + +MODELS = {"gamma": GA, "cr": CR} + +#: Column triples ``[fate, duration, censoring]`` for G1 and G2 in the phase observation. +PHASE_COLS = (np.array([0, 2, 4]), np.array([1, 3, 5])) + +TIME_FLOOR = 1e-10 + + +def outcome_mass(a: float, scale: float, p_div: float, horizon: float) -> tuple[float, float]: + """Total probability each emission assigns across the outcomes of one phase. + + For a cell watched until ``horizon`` the outcomes are: transition at some + ``t <= horizon``, death at some ``t <= horizon``, or still going at the horizon. A + well-formed emission spreads exactly probability one over those. + + :return: (Gamma/Bernoulli mass, competing-risks mass) + """ + death_scale = a * scale / max(1.0 - p_div, 1e-12) + div = sp.gamma(a, scale=scale) + death = sp.expon(scale=death_scale) + + gamma_mass = p_div * div.cdf(horizon) + (1.0 - p_div) + div.sf(horizon) + + t = np.linspace(TIME_FLOOR, horizon, 200001) + cr_mass = ( + np.trapezoid(div.pdf(t) * death.sf(t), t) + + np.trapezoid(death.pdf(t) * div.sf(t), t) + + div.sf(horizon) * death.sf(horizon) + ) + return float(gamma_mass), float(cr_mass) + + +def coarse_logpdf(dist, x: np.ndarray) -> np.ndarray: + """Log likelihood of a two-phase observation with death timing coarsened away.""" + if isinstance(dist, GA): + # The Gamma/Bernoulli emission already coarsens: a death contributes log(1 - p). + return dist.logpdf(x) + + out = np.zeros(x.shape[0]) + for cols, sub in zip(PHASE_COLS, (dist.G1, dist.G2), strict=True): + xp = x[:, cols] + divided, died, censored = event_masks(xp) + t = np.clip(xp[:, 1], TIME_FLOOR, None) + div, death = sub.div_clock, sub.death_clock + + survived = divided | censored + out[survived] += div.logsf(t[survived]) + death.logsf(t[survived]) + out[divided] += div.logpdf(t[divided]) - div.logsf(t[divided]) + out[died] += np.log(max(1.0 - sub.params[0], 1e-300)) + return out + + +def build(pops: list, cls, num_states: int, mask_seed=None, frac: float = 0.25): + """Rebuild populations under emission class ``cls``, masking ``frac`` of cells. + + Masking is driven by an rng over the tree shapes alone, so a given ``mask_seed`` + hides exactly the same cells whichever emission class is used, and the two models + are scored on identical held-out sets. + + :return: (populations, held-out records of ``(lineage index, cell indices, obs)``) + """ + E = [cls() for _ in range(num_states)] + out, held = [], [] + rng = np.random.default_rng(mask_seed) if mask_seed is not None else None + + for pop in pops: + trees, hidden = [], [] + for li, lin in enumerate(pop): + obs = lin.obs.copy() + if rng is not None: + m = rng.random(obs.shape[0]) < frac + hidden.append((li, np.nonzero(m)[0], obs[m].copy())) + # Negating an observation is how this package marks it hidden. + obs[m] *= -1.0 + trees.append(LineageTree(lin.tree, E, obs=obs, states=lin.states)) + out.append(trees) + held.append(hidden) + return out, held + + +def _state_weights(gammas, ci: int, li: int, idxs: np.ndarray) -> np.ndarray: + """Normalized posterior over states for the given held-out cells.""" + w = np.clip(gammas[ci][li][idxs, :], 1e-300, None) + return w / w.sum(axis=1, keepdims=True) + + +def heldout_LL(objs: list, gammas: list, held: list, scorer=coarse_logpdf) -> tuple[float, int]: + """Held-out log likelihood, marginalized over the fitted state posterior.""" + tot, n = 0.0, 0 + for ci, tO in enumerate(objs): + for li, idxs, true_obs in held[ci]: + if len(idxs) == 0: + continue + lp = np.stack([scorer(tO.estimate.E[s], true_obs) for s in range(tO.num_states)], axis=1) + tot += float(np.sum(logsumexp(lp + np.log(_state_weights(gammas, ci, li, idxs)), axis=1))) + n += len(idxs) + return tot, n + + +def fate_logloss(objs: list, gammas: list, held: list) -> tuple[float, int]: + """Held-out log likelihood of the binary fate alone. + + Both models emit a division probability per phase, so this compares them over an + identical observation space with no duration density involved at all. + """ + tot, n = 0.0, 0 + for ci, tO in enumerate(objs): + for li, idxs, true_obs in held[ci]: + if len(idxs) == 0: + continue + w = _state_weights(gammas, ci, li, idxs) + for ph in (0, 1): + fate = true_obs[:, ph] + known = np.isin(fate, (0.0, 1.0)) + if not np.any(known): + continue + p = np.clip([tO.estimate.E[s].params[ph] for s in range(tO.num_states)], 1e-9, 1 - 1e-9) + pf = np.where(fate[known, None] == 1.0, p[None, :], 1.0 - p[None, :]) + tot += float(np.sum(np.log(np.sum(w[known] * pf, axis=1)))) + n += int(known.sum()) + return tot, n + + +def death_time_LL(objs: list, gammas: list, held: list) -> tuple[float, int]: + """Density the competing-risks model puts on held-out death times, given a death. + + This is the piece the Gamma/Bernoulli emission cannot score at all, so it is + reported on its own rather than folded into the comparison. A positive value means + the fitted death clock is sharper than a one-per-hour reference. + """ + tot, n = 0.0, 0 + for ci, tO in enumerate(objs): + for li, idxs, true_obs in held[ci]: + if len(idxs) == 0: + continue + w = _state_weights(gammas, ci, li, idxs) + for cols, attr in zip(PHASE_COLS, ("G1", "G2"), strict=True): + xp = true_obs[:, cols] + _, died, _ = event_masks(xp) + if not np.any(died): + continue + t = np.clip(xp[died, 1], TIME_FLOOR, None) + lp = [] + for s in range(tO.num_states): + sub = getattr(tO.estimate.E[s], attr) + # Density of the death time conditional on death being the outcome. + lp.append( + sub.death_clock.logpdf(t) + sub.div_clock.logsf(t) - np.log(max(1.0 - sub.params[0], 1e-300)) + ) + tot += float(np.sum(logsumexp(np.stack(lp, axis=1) + np.log(w[died]), axis=1))) + n += int(died.sum()) + return tot, n + + +def run(pop_name: str, k_list: list[int], reps: int, seed0: int = 0): + """Fit both emissions on the same masked data and emit one JSON record per rep.""" + from . import Lineage_collections as LC + + pops = getattr(LC, pop_name) + + for k in k_list: + for r in range(reps): + mask_seed = seed0 + 1000 * k + r + row: dict = {"pop": pop_name, "k": k, "rep": r} + + for mname, cls in MODELS.items(): + trees, held = build(pops, cls, k, mask_seed=mask_seed) + t0 = time.time() + objs, LL, gam = Analyze_list(trees, k, rng=np.random.default_rng(mask_seed)) + + coarse, n = heldout_LL(objs, gam, held) + fate, n_fate = fate_logloss(objs, gam, held) + entry = { + "trainLL": LL, + "coarse_heldout": coarse, + "n_heldout": n, + "fate_LL": fate, + "n_fate": n_fate, + "dof": objs[0].estimate.E[0].dof(), + "secs": time.time() - t0, + "params": [e.params.tolist() for e in objs[0].estimate.E], + } + if cls is CR: + entry["death_time_LL"], entry["n_deaths"] = death_time_LL(objs, gam, held) + row[mname] = entry + + print(json.dumps(row), flush=True) + + +if __name__ == "__main__": + run(sys.argv[1], [int(v) for v in sys.argv[2].split(",")], int(sys.argv[3])) diff --git a/lineage/states/StateDistributionCR.py b/lineage/states/StateDistributionCR.py new file mode 100644 index 000000000..cf2f75718 --- /dev/null +++ b/lineage/states/StateDistributionCR.py @@ -0,0 +1,336 @@ +"""Competing-risks state distributions. + +The Gamma/GaPhs emissions in this package treat a cell's fate (Bernoulli) and its +phase duration (Gamma) as independent observations, and discard the duration of any +cell that dies. That throws away every death time and, more subtly, mis-states the +likelihood of a time-censored cell: "no event yet at time t" is written as +P(division > t) when it should be P(division > t AND death > t). + +Here each phase instead carries two latent clocks, + + T_D ~ Gamma(a, s) division / transition + T_X ~ Gamma(a_x, s_x) death + +and we observe min(T_D, T_X) together with an indicator of which fired. The three +likelihood cases are the standard competing-risks ones: + + transition seen at t f_D(t) * S_X(t) + death seen at t f_X(t) * S_D(t) + censored at t S_D(t) * S_X(t) + +The division probability is then *derived*, P(divide) = int f_D(t) S_X(t) dt, rather +than fit as a free Bernoulli parameter, so the death fraction and the death timing are +forced to agree. With the G1 death clock pinned to a constant hazard (shape 1) this +costs no degrees of freedom relative to the Bernoulli/Gamma model. +""" + +from typing import Literal + +import numpy as np +import scipy.stats as sp +from scipy.integrate import quad +from scipy.sparse import csr_array + +from .stateCommon import censor_lineage_gamma, censor_lineage_gaphs, gamma_estimator + +# Smallest duration fed to a pdf, so that a recorded duration of exactly zero cannot +# produce a non-finite emission likelihood. +TIME_FLOOR = 1e-10 + + +def event_masks(x: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Split cells into the three competing-risks cases. + + ``x`` has the usual three per-phase columns ``[fate, duration, censoring]``, where + fate is 1 for surviving the phase, 0 for dying in it, and NaN when unknown, and the + censoring flag is 1 when the phase was seen through to its end. + + Cells whose duration is negative have been masked for cross validation, and cells + with a NaN duration never entered the phase; both are excluded everywhere. + + :return: boolean masks for (division observed, death observed, censored) + """ + fate, dur, cens = x[:, 0], x[:, 1], x[:, 2] + + valid = np.isfinite(dur) & (dur >= 0.0) + died = valid & (fate == 0.0) + divided = valid & (fate == 1.0) & (cens == 1.0) + # Everything else that we timed is censored: an unknown fate, or a cell whose phase + # we only caught part of (the root cell's G1, or a cell first seen in G2). Those + # partial durations are lower bounds on the true one, so right-censoring them is + # the conservative reading, and it is what the Bernoulli/Gamma model already did. + censored = valid & ~died & ~divided + + return divided, died, censored + + +class StateDistribution: + """One cell-cycle phase with competing division and death clocks. + + ``params`` is ``[bern_p, gamma_a, gamma_scale, death_a, death_scale]``. The first + three entries keep the meaning they have in + :class:`~lineage.states.StateDistributionGamma.StateDistribution`, so downstream + figure code that indexes them positionally continues to work; ``bern_p`` is now + derived from the two clocks rather than fit. + """ + + def __init__( + self, + gamma_a: float = 7.0, + gamma_scale: float = 4.5, + death_a: float = 1.0, + death_scale: float = 40.0, + fixed_death_shape: bool = True, + ): + """ + :param fixed_death_shape: pin the death clock's shape at 1, i.e. a constant + death hazard. True for G1, where the death times are memoryless; False for + G2, where they show a strongly increasing hazard. + """ + self.fixed_death_shape = fixed_death_shape + if fixed_death_shape: + death_a = 1.0 + self.params = np.array([0.0, gamma_a, gamma_scale, death_a, death_scale]) + self.params[0] = self.division_probability() + + # -- the two clocks ------------------------------------------------------ + + @property + def div_clock(self): + return sp.gamma(a=self.params[1], scale=self.params[2]) + + @property + 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 + + def integrand(t): + return div.pdf(t) * death.sf(t) + + upper = div.ppf(1.0 - 1e-9) + val = quad(integrand, 0.0, upper, limit=100)[0] + return float(np.clip(val, 0.0, 1.0)) + + def rvs(self, size: int, rng=None): + """Draw min(T_D, T_X) and record which clock fired.""" + rng = np.random.default_rng(rng) + t_div = rng.gamma(self.params[1], scale=self.params[2], size=size) + t_death = rng.gamma(self.params[3], scale=self.params[4], size=size) + + divided = t_div <= t_death + return divided.astype(float), np.minimum(t_div, t_death), np.ones(size) + + def dist(self, other) -> float: + """Wasserstein distance between the division clocks of two states. + + Kept on the division clock alone so that the number is comparable with the + Bernoulli/Gamma model's. + """ + assert isinstance(self, type(other)) + return float(np.absolute(self.params[1] * self.params[2] - other.params[1] * other.params[2])) + + def dof(self) -> int: + """Two division-clock parameters plus the death clock's scale, and its shape + when that is not pinned. The Bernoulli is derived, so it is not counted.""" + 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.""" + 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 + + 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]) + + assert not np.any(np.isnan(ll)) + return ll + + def estimator(self, x: np.ndarray, gammas: np.ndarray): + """Weighted M step for a single condition. + + The two clocks separate in the complete-data likelihood, so each is just a + weighted right-censored Gamma fit over every timed cell. + """ + fit_clocks(self, [x], [gammas[:, np.newaxis]], state_j=0) + + def censor_lineage_array( + self, + censor_condition: int, + tree: csr_array, + obs: np.ndarray, + states: np.ndarray, + desired_experiment_time=2e12, + ) -> tuple[csr_array, np.ndarray, np.ndarray]: + """Applies censoring to array representation directly.""" + return censor_lineage_gamma(tree, obs, states, censor_condition, desired_experiment_time) + + +class StateDistributionPhase: + """G1 and G2 phases, each with its own pair of competing clocks. + + ``params`` is ``[bern_p1, bern_p2, a1, s1, a2, s2, death_a1, death_s1, death_a2, + death_s2]``. The leading six entries match + :class:`~lineage.states.StateDistributionGaPhs.StateDistribution` exactly. + """ + + def __init__( + self, + gamma_a1: float = 7.0, + gamma_scale1: float = 3.0, + gamma_a2: float = 14.0, + gamma_scale2: float = 6.0, + death_scale1: float = 40.0, + death_a2: float = 3.0, + death_scale2: float = 20.0, + ): + # G1 deaths are memoryless, so that clock is a one-parameter exponential; G2 + # deaths have a strongly increasing hazard and need a free shape. + self.G1 = StateDistribution(gamma_a1, gamma_scale1, 1.0, death_scale1, fixed_death_shape=True) + self.G2 = StateDistribution(gamma_a2, gamma_scale2, death_a2, death_scale2, fixed_death_shape=False) + self.params = np.empty(10) + self._sync() + + def _sync(self): + """Mirror the sub-distributions' parameters into the flat ``params`` array.""" + self.params[0] = self.G1.params[0] + self.params[1] = self.G2.params[0] + self.params[2:4] = self.G1.params[1:3] + self.params[4:6] = self.G2.params[1:3] + self.params[6:8] = self.G1.params[3:5] + self.params[8:10] = self.G2.params[3:5] + + def rvs(self, size: int, rng=None): + rng = np.random.default_rng(rng) + bern_G1, gamma_G1, cens_G1 = self.G1.rvs(size, rng=rng) + bern_G2, gamma_G2, cens_G2 = self.G2.rvs(size, rng=rng) + return bern_G1, bern_G2, gamma_G1, gamma_G2, cens_G1, cens_G2 + + def dist(self, other) -> float: + assert isinstance(self, type(other)) + return self.G1.dist(other.G1) + self.G2.dist(other.G2) + + def dof(self) -> int: + return self.G1.dof() + self.G2.dof() + + def logpdf(self, x: np.ndarray) -> np.ndarray: + return self.G1.logpdf(x[:, np.array([0, 2, 4])]) + self.G2.logpdf(x[:, np.array([1, 3, 5])]) + + def estimator(self, x: np.ndarray, gammas: np.ndarray): + self.G1.estimator(x[:, np.array([0, 2, 4])], gammas) + self.G2.estimator(x[:, np.array([1, 3, 5])], gammas) + self._sync() + + def censor_lineage_array( + self, + censor_condition: int, + tree: csr_array, + obs: np.ndarray, + states: np.ndarray, + desired_experiment_time=2e12, + ) -> tuple[csr_array, np.ndarray, np.ndarray]: + return censor_lineage_gaphs(tree, obs, states, censor_condition, desired_experiment_time) + + +def exponential_estimator(obs: np.ndarray, events: np.ndarray, weights: np.ndarray, param_idx: np.ndarray, K: int): + """Weighted right-censored MLE for exponential scales, one per group. + + With a constant hazard the MLE is total time at risk over events observed, which + needs no iteration. A pseudocount keeps a group with no observed deaths finite. + """ + scales = np.empty(K) + for k in range(K): + sel = param_idx == (k + 1) + at_risk = float(np.dot(weights[sel], obs[sel])) + 1.0 + n_events = float(np.dot(weights[sel], events[sel])) + 1.0 / K + scales[k] = at_risk / n_events + return scales + + +def fit_clocks(distributions, x_list: list[np.ndarray], gammas_list: list[np.ndarray], state_j: int): + """Fit both clocks of one state, sharing each shape across conditions. + + ``distributions`` is either a single :class:`StateDistribution` or a list of them, + one per condition. The shape parameters are shared across conditions and the + scales are free, mirroring how ``atonce_estimator`` treats the Gamma model. + """ + single = not isinstance(distributions, list) + dists = [distributions] if single else distributions + K = len(x_list) + + x = np.concatenate(x_list, axis=0) + weights = np.concatenate([g[:, state_j] for g in gammas_list]) + idx = np.concatenate([np.full(g.shape[0], k + 1) for k, g in enumerate(gammas_list)]) + + divided, died, censored = event_masks(x) + timed = divided | died | censored + if not np.any(timed): + return + + t = np.clip(x[timed, 1], TIME_FLOOR, None) + w = weights[timed] + idx = idx[timed] + div_event = divided[timed].astype(float) + death_event = died[timed].astype(float) + + ref = dists[0] + + # Division clock: every timed cell contributes, as an event or as censored. + if np.any(div_event > 0.0): + x0 = np.array([ref.params[1]] + [d.params[2] for d in dists]) + out = gamma_estimator(t, div_event, w, idx, x0, phase="all") + for k, d in enumerate(dists): + d.params[1] = out[0] + d.params[2] = out[k + 1] + + # Death clock: same cells, with the roles of event and censoring swapped. + if np.any(death_event > 0.0): + if ref.fixed_death_shape: + scales = exponential_estimator(t, death_event, w, idx, K) + for k, d in enumerate(dists): + d.params[3] = 1.0 + d.params[4] = scales[k] + else: + x0 = np.array([ref.params[3]] + [d.params[4] for d in dists]) + out = gamma_estimator(t, death_event, w, idx, x0, phase="all") + for k, d in enumerate(dists): + d.params[3] = out[0] + d.params[4] = out[k + 1] + + for d in dists: + d.params[0] = d.division_probability() + + +def atonce_estimator( + all_tHMMobj: list, + x_list: list, + gammas_list: list[np.ndarray], + phase: Literal["all", "G1", "G2"], +): + """M step across several conditions at once, matching the Gamma model's interface.""" + x_list = [np.asarray(x) for x in x_list] + + for state_j in range(len(all_tHMMobj[0].estimate.E)): + emissions = [tO.estimate.E[state_j] for tO in all_tHMMobj] + + if phase == "all": + fit_clocks(emissions, x_list, gammas_list, state_j) + else: + sub = "G1" if phase == "G1" else "G2" + fit_clocks([getattr(e, sub) for e in emissions], x_list, gammas_list, state_j) + for e in emissions: + e._sync() + + +# Let BaumWelch find the right at-once estimator from the emission object itself. +StateDistribution.atonce_estimator = staticmethod(atonce_estimator) # type: ignore[attr-defined] +StateDistributionPhase.atonce_estimator = staticmethod(atonce_estimator) # type: ignore[attr-defined] diff --git a/lineage/tests/test_StateDistributionCR.py b/lineage/tests/test_StateDistributionCR.py new file mode 100644 index 000000000..37957cabe --- /dev/null +++ b/lineage/tests/test_StateDistributionCR.py @@ -0,0 +1,160 @@ +"""Tests for the competing-risks state distributions.""" + +import numpy as np +import pytest +import scipy.stats as sp + +from lineage.compare_emissions import outcome_mass +from lineage.states.StateDistributionCR import ( + StateDistribution, + StateDistributionPhase, + event_masks, +) + + +@pytest.fixture +def dist(): + return StateDistribution(gamma_a=7.0, gamma_scale=4.5, death_a=1.0, death_scale=40.0) + + +def test_event_masks_partition(): + """Every observation lands in exactly one case, or in none when it has no time.""" + x = np.array( + [ + [1.0, 20.0, 1.0], # transition observed + [0.0, 20.0, 1.0], # death observed + [np.nan, 20.0, 0.0], # time censored, fate unknown + [1.0, 20.0, 0.0], # survived the phase but only partly observed + [0.0, 20.0, 0.0], # died, but the phase was entered before we saw it + [1.0, np.nan, np.nan], # never entered the phase + [1.0, -20.0, 1.0], # hidden for cross validation + ] + ) + divided, died, censored = event_masks(x) + + assert np.array_equal(divided, [True, False, False, False, False, False, False]) + assert np.array_equal(died, [False, True, False, False, True, False, False]) + assert np.array_equal(censored, [False, False, True, True, False, False, False]) + # Mutually exclusive, and the last two rows are excluded everywhere. + assert np.all(divided.astype(int) + died.astype(int) + censored.astype(int) <= 1) + + +def test_logpdf_matches_competing_risks_by_hand(dist): + """Each case is the textbook competing-risks term.""" + div = sp.gamma(7.0, scale=4.5) + death = sp.gamma(1.0, scale=40.0) + t = 20.0 + + x = np.array( + [ + [1.0, t, 1.0], + [0.0, t, 1.0], + [np.nan, t, 0.0], + [1.0, np.nan, np.nan], + [1.0, -t, 1.0], + ] + ) + expected = [ + div.logpdf(t) + death.logsf(t), + death.logpdf(t) + div.logsf(t), + div.logsf(t) + death.logsf(t), + 0.0, + 0.0, + ] + np.testing.assert_allclose(dist.logpdf(x), expected) + + +def test_sub_densities_sum_to_one(dist): + """The divide and die branches together carry exactly probability one.""" + t = np.linspace(1e-9, 2000.0, 400001) + div, death = dist.div_clock, dist.death_clock + + p_divide = np.trapezoid(div.pdf(t) * death.sf(t), t) + p_die = np.trapezoid(death.pdf(t) * div.sf(t), t) + + assert p_divide + p_die == pytest.approx(1.0, abs=1e-4) + # params[0] is derived from the two clocks rather than fit, so it must agree. + assert dist.params[0] == pytest.approx(p_divide, abs=1e-4) + + +def test_gamma_bernoulli_mass_exceeds_one_under_censoring(): + """The Bernoulli/Gamma emission is not normalized when cells are censored. + + A death is scored as an atom with no time attached while the censored branch uses + the division survival alone, so the two overlap. The competing-risks form does not. + """ + for horizon in (4.0, 12.0, 24.0): + gamma_mass, cr_mass = outcome_mass(a=3.38, scale=5.74, p_div=0.9, horizon=horizon) + assert gamma_mass > 1.02 + assert cr_mass == pytest.approx(1.0, abs=1e-3) + + # Both are fine once nothing is censored. + gamma_mass, cr_mass = outcome_mass(a=3.38, scale=5.74, p_div=0.9, horizon=400.0) + assert gamma_mass == pytest.approx(1.0, abs=1e-3) + + +def test_estimator_recovers_parameters(dist): + """A weighted fit to data simulated from the model returns the same clocks.""" + rng = np.random.default_rng(42) + obs = np.column_stack(dist.rvs(20000, rng=rng)) + + fitted = StateDistribution(gamma_a=1.0, gamma_scale=1.0, death_a=1.0, death_scale=1.0) + fitted.estimator(obs, np.ones(obs.shape[0])) + + np.testing.assert_allclose(fitted.params[1:], dist.params[1:], rtol=0.1) + assert fitted.params[0] == pytest.approx(dist.params[0], abs=0.02) + + +def test_rvs_is_the_minimum_of_two_clocks(dist): + """The recorded duration and fate are those of whichever clock fired first.""" + rng = np.random.default_rng(0) + fate, dur, cens = dist.rvs(5000, rng=rng) + + assert np.all(np.isin(fate, (0.0, 1.0))) + assert np.all(dur > 0.0) + assert np.all(cens == 1.0) + # Cells that divided did so faster than the death-clock mean of 40 h on average, + # and the observed division fraction tracks the derived probability. + assert fate.mean() == pytest.approx(dist.params[0], abs=0.02) + + +def test_phase_death_shapes(): + """G1's death clock is pinned to a constant hazard; G2's shape is free.""" + p = StateDistributionPhase() + + assert p.G1.fixed_death_shape and p.G1.params[3] == 1.0 + assert not p.G2.fixed_death_shape + # Three parameters for G1 (two division, one death scale) and four for G2. + assert (p.G1.dof(), p.G2.dof(), p.dof()) == (3, 4, 7) + + +def test_phase_params_layout_matches_gaphs(): + """The leading six entries keep their Gamma/Bernoulli meaning for figure code.""" + p = StateDistributionPhase(gamma_a1=7.0, gamma_scale1=3.0, gamma_a2=14.0, gamma_scale2=6.0) + + assert p.params.shape == (10,) + assert (p.params[0], p.params[1]) == (p.G1.params[0], p.G2.params[0]) + np.testing.assert_allclose(p.params[2:4], [7.0, 3.0]) + np.testing.assert_allclose(p.params[4:6], [14.0, 6.0]) + np.testing.assert_allclose(p.params[6:8], p.G1.params[3:5]) + np.testing.assert_allclose(p.params[8:10], p.G2.params[3:5]) + + +def test_phase_estimator_recovers_parameters(): + """Both phases, and both clocks within each, are recovered from simulated data.""" + truth = StateDistributionPhase( + gamma_a1=7.0, + gamma_scale1=3.0, + gamma_a2=14.0, + gamma_scale2=6.0, + death_scale1=40.0, + death_a2=3.0, + death_scale2=20.0, + ) + rng = np.random.default_rng(7) + obs = np.column_stack(truth.rvs(30000, rng=rng)) + + fitted = StateDistributionPhase(2.0, 2.0, 2.0, 2.0, death_scale1=10.0, death_a2=1.0, death_scale2=10.0) + fitted.estimator(obs, np.ones(obs.shape[0])) + + np.testing.assert_allclose(fitted.params[2:], truth.params[2:], rtol=0.15) From 7b2e07a664458ba1c26c2da1934d4ca2160c58f7 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 6 Sep 2026 20:16:47 -0700 Subject: [PATCH 2/6] Add figure S18: death timing by cell-cycle phase Shows why the two death clocks are shaped differently. G1 death times sit on the exponential (Weibull shape 0.89) while G2 death times are clearly not memoryless (shape 2.72) and arrive later than divisions do. The cumulative hazard panel makes the contrast direct: G1 death runs parallel to a constant hazard, G2 death and both divisions are much steeper. --- lineage/figures/figureS18.py | 93 ++++++++++++++++++++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 lineage/figures/figureS18.py diff --git a/lineage/figures/figureS18.py b/lineage/figures/figureS18.py new file mode 100644 index 000000000..6d3c104a1 --- /dev/null +++ b/lineage/figures/figureS18.py @@ -0,0 +1,93 @@ +"""Death timing by phase, which is what sets the shape of each death clock. + +The competing-risks emission in :mod:`lineage.states.StateDistributionCR` gives G1 a +one-parameter death clock with a constant hazard and G2 a two-parameter one. This +figure is the evidence for that split: pooled across the lapatinib and gemcitabine +conditions, G1 death times are essentially memoryless while G2 death times have a +strongly increasing hazard and arrive later than divisions do. +""" + +import numpy as np +import scipy.stats as sp + +from ..Lineage_collections import AllGemcitabine, AllLapatinib +from .common import getSetup + + +def gather() -> dict[str, np.ndarray]: + """Pool per-phase event times across the lapatinib and gemcitabine conditions. + + Observation columns are ``[G1 fate, G2 fate, G1 time, G2 time, G1 cens, G2 cens]``, + with a fate of 0 for death in that phase and 1 for surviving it. The shared control + appears in both drug lists, so populations are de-duplicated by identity. + """ + seen: dict[int, np.ndarray] = {} + for drug in (AllLapatinib, AllGemcitabine): + for population in drug: + for lineage in population: + seen.setdefault(id(lineage), lineage.obs) + x = np.vstack(list(seen.values())) + + out = {} + for name, fate, time, cens in (("G1", 0, 2, 4), ("G2", 1, 3, 5)): + t = x[:, time] + timed = np.isfinite(t) & (t > 0.0) + out[f"{name} death"] = t[timed & (x[:, fate] == 0.0)] + out[f"{name} division"] = t[timed & (x[:, fate] == 1.0) & (x[:, cens] == 1.0)] + return out + + +def cumulative_hazard(t: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Nelson-Aalen style cumulative hazard for a fully observed sample. + + On log-log axes a constant hazard plots as a line of slope 1; the slope is the + Weibull shape, so anything steeper is a wear-out process. + """ + ts = np.sort(t) + n = len(ts) + at_risk = n - np.arange(n) + return ts, np.cumsum(1.0 / at_risk) + + +def makeFigure(): + """Compare the death and division timing of each cell-cycle phase.""" + data = gather() + ax, f = getSetup((9, 3), (1, 3)) + + # (a, b) survival of the death times against the best-fit exponential. + for i, phase in enumerate(("G1", "G2")): + t = data[f"{phase} death"] + ts, _ = cumulative_hazard(t) + ax[i].step(ts, 1.0 - np.arange(len(ts)) / len(ts), where="post", label="observed") + ax[i].plot(ts, sp.expon(scale=t.mean()).sf(ts), "--", label="exponential") + shape, _, scale = sp.weibull_min.fit(t, floc=0) + ax[i].plot(ts, sp.weibull_min(shape, scale=scale).sf(ts), ":", label="Weibull") + ax[i].set( + title=f"{phase} death times (n={len(t)})", + xlabel="time in phase [hr]", + ylabel="fraction not yet dead", + ) + ax[i].text( + 0.55, + 0.75, + f"Weibull shape\n{shape:.2f}", + transform=ax[i].transAxes, + fontsize=9, + ) + ax[i].legend(fontsize=7) + + # (c) cumulative hazards on log-log axes; slope is the Weibull shape. + for label, style in ( + ("G1 death", "-"), + ("G2 death", "-"), + ("G1 division", "--"), + ("G2 division", "--"), + ): + ts, H = cumulative_hazard(data[label]) + ax[2].loglog(ts, H, style, label=f"{label} (n={len(ts)})", linewidth=1) + ref = np.array([1.0, 100.0]) + ax[2].loglog(ref, ref / 60.0, color="k", linewidth=0.5, label="slope 1 (constant hazard)") + ax[2].set(title="Cumulative hazard", xlabel="time in phase [hr]", ylabel="cumulative hazard") + ax[2].legend(fontsize=6) + + return f From eb0f220563f1a18ed8d35f6f067cc3acc0b6d6eb Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 6 Sep 2026 20:18:20 -0700 Subject: [PATCH 3/6] Declare the at-once estimator hook as a class attribute Patching it onto the classes after definition left ty unable to see it. --- lineage/states/StateDistributionCR.py | 191 +++++++++++++------------- 1 file changed, 96 insertions(+), 95 deletions(-) diff --git a/lineage/states/StateDistributionCR.py b/lineage/states/StateDistributionCR.py index cf2f75718..19e340a5e 100644 --- a/lineage/states/StateDistributionCR.py +++ b/lineage/states/StateDistributionCR.py @@ -64,6 +64,96 @@ def event_masks(x: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: return divided, died, censored +def exponential_estimator(obs: np.ndarray, events: np.ndarray, weights: np.ndarray, param_idx: np.ndarray, K: int): + """Weighted right-censored MLE for exponential scales, one per group. + + With a constant hazard the MLE is total time at risk over events observed, which + needs no iteration. A pseudocount keeps a group with no observed deaths finite. + """ + scales = np.empty(K) + for k in range(K): + sel = param_idx == (k + 1) + at_risk = float(np.dot(weights[sel], obs[sel])) + 1.0 + n_events = float(np.dot(weights[sel], events[sel])) + 1.0 / K + scales[k] = at_risk / n_events + return scales + + +def fit_clocks(distributions, x_list: list[np.ndarray], gammas_list: list[np.ndarray], state_j: int): + """Fit both clocks of one state, sharing each shape across conditions. + + ``distributions`` is either a single :class:`StateDistribution` or a list of them, + one per condition. The shape parameters are shared across conditions and the + scales are free, mirroring how ``atonce_estimator`` treats the Gamma model. + """ + single = not isinstance(distributions, list) + dists = [distributions] if single else distributions + K = len(x_list) + + x = np.concatenate(x_list, axis=0) + weights = np.concatenate([g[:, state_j] for g in gammas_list]) + idx = np.concatenate([np.full(g.shape[0], k + 1) for k, g in enumerate(gammas_list)]) + + divided, died, censored = event_masks(x) + timed = divided | died | censored + if not np.any(timed): + return + + t = np.clip(x[timed, 1], TIME_FLOOR, None) + w = weights[timed] + idx = idx[timed] + div_event = divided[timed].astype(float) + death_event = died[timed].astype(float) + + ref = dists[0] + + # Division clock: every timed cell contributes, as an event or as censored. + if np.any(div_event > 0.0): + x0 = np.array([ref.params[1]] + [d.params[2] for d in dists]) + out = gamma_estimator(t, div_event, w, idx, x0, phase="all") + for k, d in enumerate(dists): + d.params[1] = out[0] + d.params[2] = out[k + 1] + + # Death clock: same cells, with the roles of event and censoring swapped. + if np.any(death_event > 0.0): + if ref.fixed_death_shape: + scales = exponential_estimator(t, death_event, w, idx, K) + for k, d in enumerate(dists): + d.params[3] = 1.0 + d.params[4] = scales[k] + else: + x0 = np.array([ref.params[3]] + [d.params[4] for d in dists]) + out = gamma_estimator(t, death_event, w, idx, x0, phase="all") + for k, d in enumerate(dists): + d.params[3] = out[0] + d.params[4] = out[k + 1] + + for d in dists: + d.params[0] = d.division_probability() + + +def atonce_estimator( + all_tHMMobj: list, + x_list: list, + gammas_list: list[np.ndarray], + phase: Literal["all", "G1", "G2"], +): + """M step across several conditions at once, matching the Gamma model's interface.""" + x_list = [np.asarray(x) for x in x_list] + + for state_j in range(len(all_tHMMobj[0].estimate.E)): + emissions = [tO.estimate.E[state_j] for tO in all_tHMMobj] + + if phase == "all": + fit_clocks(emissions, x_list, gammas_list, state_j) + else: + sub = "G1" if phase == "G1" else "G2" + fit_clocks([getattr(e, sub) for e in emissions], x_list, gammas_list, state_j) + for e in emissions: + e._sync() + + class StateDistribution: """One cell-cycle phase with competing division and death clocks. @@ -74,6 +164,9 @@ class StateDistribution: derived from the two clocks rather than fit. """ + #: BaumWelch looks this up on the emission object to pick the right M step. + atonce_estimator = staticmethod(atonce_estimator) + def __init__( self, gamma_a: float = 7.0, @@ -183,6 +276,9 @@ class StateDistributionPhase: :class:`~lineage.states.StateDistributionGaPhs.StateDistribution` exactly. """ + #: BaumWelch looks this up on the emission object to pick the right M step. + atonce_estimator = staticmethod(atonce_estimator) + def __init__( self, gamma_a1: float = 7.0, @@ -239,98 +335,3 @@ def censor_lineage_array( desired_experiment_time=2e12, ) -> tuple[csr_array, np.ndarray, np.ndarray]: return censor_lineage_gaphs(tree, obs, states, censor_condition, desired_experiment_time) - - -def exponential_estimator(obs: np.ndarray, events: np.ndarray, weights: np.ndarray, param_idx: np.ndarray, K: int): - """Weighted right-censored MLE for exponential scales, one per group. - - With a constant hazard the MLE is total time at risk over events observed, which - needs no iteration. A pseudocount keeps a group with no observed deaths finite. - """ - scales = np.empty(K) - for k in range(K): - sel = param_idx == (k + 1) - at_risk = float(np.dot(weights[sel], obs[sel])) + 1.0 - n_events = float(np.dot(weights[sel], events[sel])) + 1.0 / K - scales[k] = at_risk / n_events - return scales - - -def fit_clocks(distributions, x_list: list[np.ndarray], gammas_list: list[np.ndarray], state_j: int): - """Fit both clocks of one state, sharing each shape across conditions. - - ``distributions`` is either a single :class:`StateDistribution` or a list of them, - one per condition. The shape parameters are shared across conditions and the - scales are free, mirroring how ``atonce_estimator`` treats the Gamma model. - """ - single = not isinstance(distributions, list) - dists = [distributions] if single else distributions - K = len(x_list) - - x = np.concatenate(x_list, axis=0) - weights = np.concatenate([g[:, state_j] for g in gammas_list]) - idx = np.concatenate([np.full(g.shape[0], k + 1) for k, g in enumerate(gammas_list)]) - - divided, died, censored = event_masks(x) - timed = divided | died | censored - if not np.any(timed): - return - - t = np.clip(x[timed, 1], TIME_FLOOR, None) - w = weights[timed] - idx = idx[timed] - div_event = divided[timed].astype(float) - death_event = died[timed].astype(float) - - ref = dists[0] - - # Division clock: every timed cell contributes, as an event or as censored. - if np.any(div_event > 0.0): - x0 = np.array([ref.params[1]] + [d.params[2] for d in dists]) - out = gamma_estimator(t, div_event, w, idx, x0, phase="all") - for k, d in enumerate(dists): - d.params[1] = out[0] - d.params[2] = out[k + 1] - - # Death clock: same cells, with the roles of event and censoring swapped. - if np.any(death_event > 0.0): - if ref.fixed_death_shape: - scales = exponential_estimator(t, death_event, w, idx, K) - for k, d in enumerate(dists): - d.params[3] = 1.0 - d.params[4] = scales[k] - else: - x0 = np.array([ref.params[3]] + [d.params[4] for d in dists]) - out = gamma_estimator(t, death_event, w, idx, x0, phase="all") - for k, d in enumerate(dists): - d.params[3] = out[0] - d.params[4] = out[k + 1] - - for d in dists: - d.params[0] = d.division_probability() - - -def atonce_estimator( - all_tHMMobj: list, - x_list: list, - gammas_list: list[np.ndarray], - phase: Literal["all", "G1", "G2"], -): - """M step across several conditions at once, matching the Gamma model's interface.""" - x_list = [np.asarray(x) for x in x_list] - - for state_j in range(len(all_tHMMobj[0].estimate.E)): - emissions = [tO.estimate.E[state_j] for tO in all_tHMMobj] - - if phase == "all": - fit_clocks(emissions, x_list, gammas_list, state_j) - else: - sub = "G1" if phase == "G1" else "G2" - fit_clocks([getattr(e, sub) for e in emissions], x_list, gammas_list, state_j) - for e in emissions: - e._sync() - - -# Let BaumWelch find the right at-once estimator from the emission object itself. -StateDistribution.atonce_estimator = staticmethod(atonce_estimator) # type: ignore[attr-defined] -StateDistributionPhase.atonce_estimator = staticmethod(atonce_estimator) # type: ignore[attr-defined] From 1c5cab57f8f7d6a1f9c05c9df8a1350391d7386e Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 6 Sep 2026 20:22:14 -0700 Subject: [PATCH 4/6] Estimate the Bernoulli from the same filtered cells as the Gamma fit StateDistribution.estimator built filtered copies of every observation array to drop cross-validation-masked cells, then passed the *unfiltered* arrays to bern_estimator. Masking negates a cell's whole observation, so a hidden dividing cell reads as -1: finite, and so counted in the denominator, but never equal to 1, and so absent from the numerator. With 25% of cells hidden that pulls a true 0.90 division probability down to 0.67. Only the one-condition M step was affected; atonce_estimator already filters before estimating. This is independent of the competing-risks work and can be taken on its own. --- lineage/states/StateDistributionGamma.py | 2 +- lineage/tests/test_StateDistribution.py | 30 ++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/lineage/states/StateDistributionGamma.py b/lineage/states/StateDistributionGamma.py index 572e83b7a..050934766 100644 --- a/lineage/states/StateDistributionGamma.py +++ b/lineage/states/StateDistributionGamma.py @@ -99,7 +99,7 @@ def estimator(self, x: np.ndarray, gammas: np.ndarray): g_mask = np.logical_and(np.isfinite(γ_obs_), bern_obs_.astype("bool")) assert np.sum(g_mask) > 0, "All the cells are eliminated from the Gamma estimator." - self.params[0] = bern_estimator(bern_obs, gammas) + self.params[0] = bern_estimator(bern_obs_, gammas_) param_idx = np.ones((gammas_[g_mask].size), dtype=int) self.params[1], self.params[2] = gamma_estimator( diff --git a/lineage/tests/test_StateDistribution.py b/lineage/tests/test_StateDistribution.py index ffff86d22..253552d3b 100644 --- a/lineage/tests/test_StateDistribution.py +++ b/lineage/tests/test_StateDistribution.py @@ -114,3 +114,33 @@ def test_self_dist_zero(dist): """Test that the distance from a distribution to itself is zero.""" dd = dist() assert dd.dist(dd) == 0.0 + + +def test_bern_estimator_ignores_masked_cells(): + """Cross-validation masking must not drag the Bernoulli estimate down. + + hide_observation marks a cell by negating its whole observation, so a hidden + dividing cell reads as -1: finite, and therefore counted in the denominator, but + never equal to 1 and so missing from the numerator. Estimating from the same + filtered arrays the Gamma fit uses keeps it out of both. + """ + rng = np.random.default_rng(0) + n, p_true = 20000, 0.9 + x = np.column_stack( + [ + rng.binomial(1, p_true, n).astype(float), + rng.gamma(7.0, 4.5, n), + np.ones(n), + ] + ) + + unmasked = StateDistribution() + unmasked.estimator(x, np.ones(n)) + + masked_x = x.copy() + masked_x[rng.random(n) < 0.25] *= -1.0 + masked = StateDistribution() + masked.estimator(masked_x, np.ones(n)) + + assert unmasked.params[0] == pytest.approx(p_true, abs=0.01) + assert masked.params[0] == pytest.approx(p_true, abs=0.01) From 0d85037c868e1ff61157233d77b50aa35d17505d Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 6 Sep 2026 20:23:58 -0700 Subject: [PATCH 5/6] Document the competing-risks state distribution The state-distribution guide teaches the Bernoulli/Gamma form without noting that it discards death times and leaves the emission unnormalized under censoring. Add a section covering the competing-risks alternative, why the division probability becomes derived rather than fit, and why each phase gets a differently shaped death clock. --- docs/stateDistributions.rst | 48 +++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/docs/stateDistributions.rst b/docs/stateDistributions.rst index ef9734e9e..0b119ab0d 100644 --- a/docs/stateDistributions.rst +++ b/docs/stateDistributions.rst @@ -287,3 +287,51 @@ transition matrix and state parameters. print(" estimated state:", tHMMobj.estimate.E[state]) print("original parameters given for state:", E[state]) print("\n") + +Competing risks: tying lifetime to fate +--------------------------------------- + +The Bernoulli/Gamma example above treats a cell's fate and its lifetime as two +independent observations, and it discards the lifetime of any cell that dies -- +that is what ``ll[x[:, 0] == 0] = 0.0`` does. It also scores a time-censored cell +with ``logsf`` of the division clock alone, which states that the cell has not yet +divided but says nothing about it not having died. Those two choices together leave +the emission unnormalized: summed over the outcomes of a cell watched to a finite +horizon, its total probability exceeds one whenever there is censoring. + +``lineage.states.StateDistributionCR`` shows the alternative. Each phase carries two +latent clocks -- a division clock ``T_D`` and a death clock ``T_X`` -- and what we +observe is ``min(T_D, T_X)`` together with which one fired. The likelihood has the +three standard competing-risks cases: + +.. code:: ipython3 + + # transition seen at t: f_D(t) * S_X(t) + # death seen at t: f_X(t) * S_D(t) + # censored at t: S_D(t) * S_X(t) + + divided, died, censored = event_masks(x) + timed = divided | died | censored + + ll = np.zeros(x.shape[0]) + # every timed cell survived both clocks up to t ... + ll[timed] += div.logsf(t[timed]) + death.logsf(t[timed]) + # ... and whichever clock fired swaps its survival term for a density + ll[divided] += div.logpdf(t[divided]) - div.logsf(t[divided]) + ll[died] += death.logpdf(t[died]) - death.logsf(t[died]) + +Two things follow. Death times now carry information rather than being thrown away, +and the division probability stops being a free parameter: it is derived as +``P(divide) = int f_D(t) S_X(t) dt``, so the fraction of cells that die and the times +at which they die are forced to agree with each other. + +The estimator is no harder than before. In the complete-data likelihood the two +clocks separate, so each is an independently weighted right-censored fit over exactly +the same cells -- once with "transition observed" as the event indicator, once with +"death observed". Both reuse ``gamma_estimator`` from ``stateCommon``. + +Choosing the death clock is an empirical question, and the answer differs by phase in +the Heiser lab data (see ``lineage/figures/figureS18.py``). G1 death times are +memoryless, so that clock is a one-parameter exponential and the phase costs no more +degrees of freedom than the Bernoulli/Gamma form did. G2 death times have a strongly +increasing hazard, so that clock keeps a free shape. From 9d26c71fa682b4a0ac363f9281161d1a2ccbbad2 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 6 Sep 2026 20:27:14 -0700 Subject: [PATCH 6/6] Test the competing-risks emission through the package cross validation Covers rand_init, array censoring, hide_observation and the at-once M step in one pass, and confirms two states beat one on held-out likelihood. --- lineage/tests/test_StateDistributionCR.py | 32 +++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/lineage/tests/test_StateDistributionCR.py b/lineage/tests/test_StateDistributionCR.py index 37957cabe..e0d354b7c 100644 --- a/lineage/tests/test_StateDistributionCR.py +++ b/lineage/tests/test_StateDistributionCR.py @@ -158,3 +158,35 @@ def test_phase_estimator_recovers_parameters(): fitted.estimator(obs, np.ones(obs.shape[0])) np.testing.assert_allclose(fitted.params[2:], truth.params[2:], rtol=0.15) + + +def test_works_through_crossval(): + """The competing-risks emission drops into the package's own cross validation. + + Two well-separated states should beat one on held-out likelihood, and the run + exercises rand_init, censoring, hide_observation and the at-once M step. + """ + from lineage.BaumWelch import calculate_stationary + from lineage.crossval import crossval, hide_observation + from lineage.LineageTree import LineageTree + + T = np.array([[0.9, 0.1], [0.1, 0.9]]) + E = [ + StateDistributionPhase(8.0, 7.0, 4.0, 2.0, death_scale1=400.0, death_a2=3.0, death_scale2=200.0), + StateDistributionPhase(6.0, 4.0, 3.0, 5.0, death_scale1=20.0, death_a2=3.0, death_scale2=12.0), + ] + rng = np.random.default_rng(3) + + populations = [ + [ + LineageTree.rand_init( + calculate_stationary(T), T, E, 31, censor_condition=3, desired_experiment_time=150, rng=rng + ) + for _ in range(20) + ] + for _ in range(3) + ] + train = [hide_observation(pop, 0.25, rng=rng) for pop in populations] + + ll = crossval(train, np.arange(1, 3), rng=rng) + assert ll[0] < ll[1]