@@ -10,8 +10,11 @@
 from sklearn.metrics import roc_curve
 from statsmodels.duration.survfunc import SurvfuncRight
 
+from ..BaumWelch import do_E_step
 from ..palbociclib_analysis import CONDITIONS, OUTPUT
 from ..palbociclib_loader import load_lineages
+from ..states.CensoredWeibullGaussian import StateDistribution
+from ..tHMM import tHMM
 from .common import getSetup, subplotLabel
 
 STATE_NAMES = ("sensitive", "tolerant")
@@ -24,37 +27,52 @@
 }
 
 
+def posteriors(res: dict, pops: list) -> list[np.ndarray]:
+    """Per-cell state posteriors under the saved per-dose fit (one E step, no refitting)."""
+    E = [StateDistribution(*p) for p in res["emissions"]]
+    out = []
+    for pop, T, pi in zip(pops, res["point"]["T"], res["pi"], strict=True):
+        tO = tHMM(pop, num_states=len(E), fpi=np.array(pi), fT=np.array(T), fE=E)
+        out.append(np.vstack(do_E_step(tO)[3]))
+    return out
+
+
 def makeFigure():
     with open(OUTPUT) as f:
         res = json.load(f)
     pops = [load_lineages(c) for c in CONDITIONS]
     obs = [np.vstack([lin.obs for lin in pop]) for pop in pops]
+    gammas = posteriors(res, pops)
     E = np.array(res["emissions"])  # [mu, sigma, kappa, lam] per state
 
     ax, f = getSetup((12, 7), (2, 3))
 
-    # (a) CDK2 activation rate with the fitted state mixture for each condition.
+    # (a) CDK2 activation rate with the fitted state mixture for each condition, weighted
+    # by how many of the measured cells the model puts in each state.
     grid = np.linspace(-0.3, 1.2, 300)
     for d, (cond, o) in enumerate(zip(CONDITIONS, obs, strict=True)):
-        x = o[np.isfinite(o[:, 0]), 0]
-        ax[0].hist(x, bins=np.linspace(-0.3, 1.2, 60), density=True, alpha=0.35, color=COND_COLORS[d], label=cond)
-        pi = res["pi"][d]
-        mix = sum(pi[k] * sp.norm.pdf(grid, E[k, 0], E[k, 1]) for k in range(len(pi)))
+        m = np.isfinite(o[:, 0])
+        ax[0].hist(o[m, 0], bins=np.linspace(-0.3, 1.2, 60), density=True, alpha=0.35, color=COND_COLORS[d], label=cond)
+        w = gammas[d][m].mean(axis=0)
+        mix = sum(w[k] * sp.norm.pdf(grid, E[k, 0], E[k, 1]) for k in range(len(w)))
         ax[0].plot(grid, mix, color=COND_COLORS[d])
     ax[0].set_xlabel("max CDK2 activation rate in G1 (1/h)")
     ax[0].set_ylabel("density")
     ax[0].legend()
 
-    # (b) Kaplan-Meier of lifetimes vs the fitted Weibull mixture.
+    # (b) Kaplan-Meier of lifetimes vs the fitted Weibull mixture over the same cells.
+    # Only cells born into drug: lineage roots are in the data because they divided, so
+    # their lifetimes are selected on the outcome.
     tgrid = np.linspace(0, 25, 200)
     for d, (cond, o) in enumerate(zip(CONDITIONS, obs, strict=True)):
-        m = np.isfinite(o[:, 1])
+        not_root = np.concatenate([np.arange(len(lin)) > 0 for lin in pops[d]])
+        m = np.isfinite(o[:, 1]) & not_root
         km = SurvfuncRight(o[m, 1], o[m, 2])
         ax[1].step(km.surv_times, km.surv_prob, where="post", color=COND_COLORS[d], label=f"{cond} (KM)")
-        pi = res["pi"][d]
-        S = sum(pi[k] * np.exp(-((tgrid / E[k, 3]) ** E[k, 2])) for k in range(len(pi)))
+        w = gammas[d][m].mean(axis=0)
+        S = sum(w[k] * np.exp(-((tgrid / E[k, 3]) ** E[k, 2])) for k in range(len(w)))
         ax[1].plot(tgrid, S, "--", color=COND_COLORS[d], label=f"{cond} (tHMM)")
-    ax[1].set_xlabel("time since birth (h)")
+    ax[1].set_xlabel("time since birth into drug (h)")
     ax[1].set_ylabel("fraction undivided")
     ax[1].set_ylim(0, 1.02)
     ax[1].legend(fontsize=7)
@@ -76,14 +94,14 @@ def makeFigure():
     ax[2].set_ylim(0, 1)
     ax[2].legend(fontsize=7)
 
-    # (d) Heritability beyond state frequency: T_22 against the stationary tolerant
-    # fraction (the value T_22 takes with no memory), and the memory eigenvalue.
-    pi_tol = np.array([p[-1] for p in res["pi"]])
+    # (d) Heritability beyond state frequency: T_22 against its value in the fitted
+    # no-memory model (identical rows), and the memory eigenvalue.
+    no_memory = np.array([T_ind[1][1] for T_ind in res["independent_T"]])
     lam = np.array(res["point"]["memory_eigenvalue"])
     lam_lo, lam_hi = (np.array(a) for a in res["bootstrap"]["memory_eigenvalue"])
     xpos = np.arange(len(CONDITIONS))
     ax[3].bar(xpos - width / 2, T[:, 1, 1], width, label="$T_{22}$", color="#dd8452")
-    ax[3].bar(xpos + width / 2, pi_tol, width, label=r"$\pi_2$ (no-memory $T_{22}$)", color="#bbbbbb")
+    ax[3].bar(xpos + width / 2, no_memory, width, label=r"$T_{22}$, no-memory model", color="#bbbbbb")
     ax[3].errorbar(xpos, lam, yerr=np.vstack([lam - lam_lo, lam_hi - lam]), fmt="ko", capsize=3, label=r"$\lambda_2$")
     ax[3].set_xticks(xpos, list(CONDITIONS))
     ax[3].set_ylim(0, 1)
@@ -123,6 +141,7 @@ def makeFigure():
     ax[5].set_xlabel("number of states")
     ax[5].set_ylabel("BIC")
     ax[5].set_xticks([s["states"] for s in sel])
+    ax[5].margins(x=0.08, y=0.08)
 
     subplotLabel(ax)
     return f
@@ -19,6 +19,7 @@
 null directly with a likelihood-ratio test.
 """
 
+from collections.abc import Callable
 from concurrent.futures import ProcessPoolExecutor
 from copy import deepcopy
 from dataclasses import dataclass, field
@@ -118,6 +119,11 @@ def _random_fit(pop_list, num_states, shared_T, independent_T, estimate_pi, seed
     return objs, LL
 
 
+def _warm_fit(pop_list, template, shared_T, independent_T, estimate_pi, seed):
+    objs, LL, _ = refit(pop_list, template, shared_T, independent_T, estimate_pi, rng=seed)
+    return objs, LL
+
+
 def fit_best(
     pop_list: list,
     num_states: int,
@@ -136,19 +142,21 @@ def fit_best(
     time course they are cells from before the drug -- and the stationary tie also
     breaks the monotonicity of EM, which the likelihood-ratio tests rely on.
 
-    :param n_jobs: run the random restarts in this many processes
+    :param n_jobs: run the random restarts and warm starts in this many processes
     """
     rng = np.random.default_rng(rng)
-    seeds = rng.integers(2**32, size=n_starts)
+    seeds = rng.integers(2**32, size=n_starts + len(starts))
     args = (pop_list, num_states, shared_T, independent_T, estimate_pi)
-    if n_jobs > 1:
-        with ProcessPoolExecutor(min(n_jobs, n_starts)) as exe:
-            fits = list(exe.map(_random_fit, *zip(*[(*args, s) for s in seeds], strict=True)))
+    jobs: list[tuple[Callable, tuple]] = [(_random_fit, (*args, s)) for s in seeds[:n_starts]]
+    jobs += [
+        (_warm_fit, (pop_list, st, shared_T, independent_T, estimate_pi, s))
+        for st, s in zip(starts, seeds[n_starts:], strict=True)
+    ]
+    if n_jobs > 1 and len(jobs) > 1:
+        with ProcessPoolExecutor(min(n_jobs, len(jobs))) as exe:
+            fits = [f.result() for f in [exe.submit(fn, *a) for fn, a in jobs]]
     else:
-        fits = [_random_fit(*args, s) for s in seeds]
-    for s in starts:
-        objs, LL, _ = refit(pop_list, s, shared_T, independent_T, estimate_pi, rng=rng)
-        fits.append((objs, LL))
+        fits = [fn(*a) for fn, a in jobs]
     return max(fits, key=lambda f: f[1])
 
 
@@ -207,11 +215,22 @@ def dose_sweep(
     shared = fit_best(pops_by_dose, K, shared_T=True, **kw)
     indep = fit_best(pops_by_dose, K, shared_T=False, independent_T=True, **kw)
     per_dose = fit_best(pops_by_dose, K, shared_T=False, starts=(shared[0], indep[0]), **kw)
-    shared = fit_best(pops_by_dose, K, shared_T=True, starts=(shared[0], per_dose[0]), n_starts=0, rng=rng)
+    shared = fit_best(
+        pops_by_dose, K, shared_T=True, starts=(shared[0], per_dose[0]), n_starts=0, n_jobs=n_jobs, rng=rng
+    )
     indep = fit_best(
-        pops_by_dose, K, shared_T=False, independent_T=True, starts=(indep[0], per_dose[0]), n_starts=0, rng=rng
+        pops_by_dose,
+        K,
+        shared_T=False,
+        independent_T=True,
+        starts=(indep[0], per_dose[0]),
+        n_starts=0,
+        n_jobs=n_jobs,
+        rng=rng,
+    )
+    per_dose = fit_best(
+        pops_by_dose, K, shared_T=False, starts=(per_dose[0], shared[0], indep[0]), n_starts=0, n_jobs=n_jobs, rng=rng
     )
-    per_dose = fit_best(pops_by_dose, K, shared_T=False, starts=(per_dose[0], shared[0], indep[0]), n_starts=0, rng=rng)
 
     for fit in (per_dose, shared, indep):
         order_by_lifetime(fit[0])
 
@@ -1,7 +1,7 @@
 """Issue #1016: heritability of palbociclib escape in Spencer-lab MCF10A lineages.
 
 Run as ``OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m lineage.palbociclib_analysis
-[n_boot]``. The fits run in parallel worker processes, so BLAS's own threads should be
+[n_boot | sensitivity]``. The fits run in parallel worker processes, so BLAS's own threads should be
 pinned to one, or the workers oversubscribe the machine. Results are written to
 ``output/palbociclib_analysis.json`` and drawn by :mod:`lineage.figures.figure22`.
 
@@ -86,6 +86,10 @@ def percentile_ci(samples: np.ndarray) -> list:
     return [np.asarray(lo).tolist(), np.asarray(hi).tolist()]
 
 
+def log(msg: str, t0: float):
+    print(f"[{time.time() - t0:7.0f} s] {msg}", file=sys.stderr, flush=True)
+
+
 def run(n_boot: int = 200, n_workers: int = 16, seed: int = 0) -> dict:
     rng = np.random.default_rng(seed)
     pops = [load_lineages(c) for c in CONDITIONS]
@@ -102,8 +106,10 @@ def run(n_boot: int = 200, n_workers: int = 16, seed: int = 0) -> dict:
 
     t0 = time.time()
     res["state_selection"] = select_states(pops, n_jobs=n_workers, rng=rng)
+    log("state selection done", t0)
 
     sweep = dose_sweep(pops, list(DOSES_NM), num_states=2, n_starts=12, n_jobs=n_workers, rng=rng)
+    log("dose sweep done", t0)
     res["LL"] = {k: float(v) for k, v in sweep.LL.items()}
     res["lrt"] = sweep.lrt
     res["emissions"] = [e.params.tolist() for e in sweep.per_dose[0].estimate.E]
@@ -121,6 +127,7 @@ def run(n_boot: int = 200, n_workers: int = 16, seed: int = 0) -> dict:
     with ProcessPoolExecutor(n_workers) as exe:
         futs = [exe.submit(bootstrap_transitions, sweep, pops, len(c), s) for c, s in zip(chunks, seeds, strict=True)]
         boots = np.concatenate([f.result() for f in futs if f.result().size])
+    log("bootstrap done", t0)
     summaries = [summarize_T(b) for b in boots]
     boot: dict = {
         key: percentile_ci(np.array([s[key] for s in summaries], dtype=float))
@@ -142,6 +149,7 @@ def run(n_boot: int = 200, n_workers: int = 16, seed: int = 0) -> dict:
     res["n_boot"] = int(boots.shape[0])
 
     cv = cross_validated_scores(sweep, pops, horizon=HORIZON_H, n_folds=5, escape_col=ESCAPE_COL, rng=rng)
+    log("cross-validated forecasts done", t0)
     res["auc"] = auc_table(cv, n_boot=1000, rng=rng)
     pairs = cv["pairs"]
     res["pairs"] = {
@@ -159,10 +167,63 @@ def run(n_boot: int = 200, n_workers: int = 16, seed: int = 0) -> dict:
     return res
 
 
+SENSITIVITY_OUTPUT = os.path.join("output", "palbociclib_sensitivity.json")
+
+
+def sensitivity(n_workers: int = 16, seed: int = 1) -> dict:
+    """Refit the dose sweep without the roots' lifetimes, which are selected on dividing.
+
+    Also reports how well each fit reproduces division among cells born into drug, the
+    cells not selected on their outcome: the observed fraction divided by 20 h (Kaplan-Meier)
+    against the fraction the fitted states predict.
+    """
+    from statsmodels.duration.survfunc import SurvfuncRight
+
+    from .BaumWelch import do_E_step
+    from .tHMM import tHMM
+
+    rng = np.random.default_rng(seed)
+    out: dict = {}
+    for label, keep in (("with_root_lifetimes", True), ("without_root_lifetimes", False)):
+        pops = [load_lineages(c, root_lifetimes=keep) for c in CONDITIONS]
+        sweep = dose_sweep(pops, list(DOSES_NM), num_states=2, n_starts=12, n_jobs=n_workers, rng=rng)
+        E = sweep.per_dose[0].estimate.E
+        entry = {
+            "LL": {k: float(v) for k, v in sweep.LL.items()},
+            "lrt": sweep.lrt,
+            "emissions": [e.params.tolist() for e in E],
+            "pi": [tO.estimate.pi.tolist() for tO in sweep.per_dose],
+            "point": summarize_T(sweep.T),
+            "divided_by_20h": [],
+        }
+        # Always score division on the full data, whichever lineages were fit.
+        full = [load_lineages(c) for c in CONDITIONS]
+        for pop, tO in zip(full, sweep.per_dose, strict=True):
+            fixed = tHMM(pop, 2, fpi=tO.estimate.pi, fT=tO.estimate.T, fE=E)
+            gam = np.vstack(do_E_step(fixed)[3])
+            obs = np.vstack([lin.obs for lin in pop])
+            m = np.isfinite(obs[:, 1]) & np.concatenate([np.arange(len(lin)) > 0 for lin in pop])
+            km = SurvfuncRight(obs[m, 1], obs[m, 2])
+            w = gam[m].mean(axis=0)
+            S20 = sum(w[k] * np.exp(-((20.0 / e.params[3]) ** e.params[2])) for k, e in enumerate(E))
+            entry["divided_by_20h"].append(
+                {"observed": float(1.0 - np.interp(20.0, km.surv_times, km.surv_prob)), "model": float(1.0 - S20)}
+            )
+        out[label] = entry
+    return out
+
+
 if __name__ == "__main__":
-    n_boot = int(sys.argv[1]) if len(sys.argv) > 1 else 200
-    out = run(n_boot=n_boot)
-    os.makedirs(os.path.dirname(OUTPUT), exist_ok=True)
-    with open(OUTPUT, "w") as f:
-        json.dump(out, f, indent=1, default=float)
-    print(json.dumps({k: v for k, v in out.items() if not k.startswith("roc_")}, indent=1, default=float))
+    if len(sys.argv) > 1 and sys.argv[1] == "sensitivity":
+        out = sensitivity()
+        os.makedirs(os.path.dirname(SENSITIVITY_OUTPUT), exist_ok=True)
+        with open(SENSITIVITY_OUTPUT, "w") as f:
+            json.dump(out, f, indent=1, default=float)
+        print(json.dumps(out, indent=1, default=float))
+    else:
+        n_boot = int(sys.argv[1]) if len(sys.argv) > 1 else 200
+        out = run(n_boot=n_boot)
+        os.makedirs(os.path.dirname(OUTPUT), exist_ok=True)
+        with open(OUTPUT, "w") as f:
+            json.dump(out, f, indent=1, default=float)
+        print(json.dumps({k: v for k, v in out.items() if not k.startswith("roc_")}, indent=1, default=float))
@@ -189,7 +189,9 @@ def cell_obs(df: pd.DataFrame) -> np.ndarray:
     )
 
 
-def load_lineages(condition: str, table: pd.DataFrame | None = None, E=None) -> list[LineageTree]:
+def load_lineages(
+    condition: str, table: pd.DataFrame | None = None, E=None, root_lifetimes: bool = True
+) -> list[LineageTree]:
     """Lineages of cells born after drug addition in one condition.
 
     Each lineage is rooted at a cell that divided after the drug went on; its daughters,
@@ -202,6 +204,10 @@ def load_lineages(condition: str, table: pd.DataFrame | None = None, E=None) ->
     :param condition: ``"control"``, ``"palbociclib"``, ``"MEKi"``, or ``"Nutlin"``
     :param table: a table from :func:`build_cell_table`; defaults to :data:`CELL_TABLE`
     :param E: emissions to attach to the lineages (only used as a template for fitting)
+    :param root_lifetimes: keep the roots' lifetimes. Roots are only in the data because
+        they divided before the movie ended, so their lifetimes are selected on the
+        outcome; pass False to drop them (keeping the roots' biosensor readings) as a
+        sensitivity check on that selection.
     """
     if table is None:
         table = pd.read_csv(CELL_TABLE)
@@ -236,5 +242,7 @@ def load_lineages(condition: str, table: pd.DataFrame | None = None, E=None) ->
             if root not in movie.index:
                 # A mother with no usable sensor frames at all is kept for her topology only.
                 obs[0, :] = np.nan
+            if not root_lifetimes:
+                obs[0, 1:3] = np.nan
             lineages.append(LineageTree(tree, E, obs=obs))
     return lineages
@@ -141,3 +141,12 @@ def test_shipped_table_loads():
             assert np.all(parents < daughters)
             # Two daughters, or one when the sister never had a usable sensor frame.
             assert np.all(np.isin(np.bincount(parents, minlength=len(lin))[parents], (1, 2)))
+
+
+def test_load_lineages_without_root_lifetimes(movie_dir):
+    table = build_cell_table(str(movie_dir))
+    [lin] = load_lineages("palbociclib", table, root_lifetimes=False)
+    [full] = load_lineages("palbociclib", table)
+    assert np.all(np.isnan(lin.obs[0, 1:3]))
+    np.testing.assert_array_equal(lin.obs[0, 0], full.obs[0, 0])
+    np.testing.assert_array_equal(lin.obs[1:], full.obs[1:])