|
| 1 | +"""Heritability of the CDK2 state under timed MEK or ERK inhibition (S-BSST314 Figure 1). |
| 2 | +
|
| 3 | +Run as ``OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m lineage.mitogen_analysis |
| 4 | +[n_boot]``. For each drug, the six exposures (none, 1, 3, 6, 9 h, left on) are fit as a |
| 5 | +dose series by :func:`.heritability.dose_sweep`, with lineage roots' lifetimes truncated |
| 6 | +to the window in which they were selected (see :func:`.palbociclib_loader.build_lineages`). |
| 7 | +Results are written to ``output/mitogen_analysis.json`` and drawn by |
| 8 | +:mod:`lineage.figures.figure23`. |
| 9 | +
|
| 10 | +See :mod:`lineage.mitogen_loader` for the data. |
| 11 | +""" |
| 12 | + |
| 13 | +import json |
| 14 | +import os |
| 15 | +import sys |
| 16 | +import time |
| 17 | +from concurrent.futures import ThreadPoolExecutor |
| 18 | + |
| 19 | +import numpy as np |
| 20 | + |
| 21 | +from .heritability import dose_sweep |
| 22 | +from .mitogen_loader import DRUGS, PULSES_H, generation_depths, load_lineages |
| 23 | +from .palbociclib_analysis import bootstrap_summary, division_calibration, log, summarize_T |
| 24 | +from .palbociclib_loader import ROOT_MODES |
| 25 | +from .states.CensoredWeibullGaussian import StateDistribution |
| 26 | + |
| 27 | +OUTPUT = os.path.join("output", "mitogen_analysis.json") |
| 28 | + |
| 29 | +#: Lineages kept per condition. The movies hold 2,000-18,000 lineages per condition, and |
| 30 | +#: EM time is dominated by per-lineage overhead, so each condition is subsampled to about |
| 31 | +#: the size of the palbociclib data (the stay-in MEKi wells have fewer and are kept whole). |
| 32 | +MAX_LINEAGES = 2500 |
| 33 | + |
| 34 | + |
| 35 | +def condition_lineages(drug: str, rng, roots: tuple[str, ...] = ROOT_MODES) -> tuple[dict[str, list[list]], list[int]]: |
| 36 | + """Per-condition lineages under each root treatment, with the same random subset of |
| 37 | + lineages for every treatment, and how many lineages each condition had to draw from.""" |
| 38 | + out: dict[str, list[list]] = {mode: [] for mode in roots} |
| 39 | + available = [] |
| 40 | + for p in PULSES_H: |
| 41 | + full = {mode: load_lineages(drug, p, roots=mode) for mode in roots} |
| 42 | + n = len(full[roots[0]]) |
| 43 | + available.append(n) |
| 44 | + keep = np.sort(rng.choice(n, size=min(n, MAX_LINEAGES), replace=False)) |
| 45 | + for mode in roots: |
| 46 | + out[mode].append([full[mode][i] for i in keep]) |
| 47 | + return out, available |
| 48 | + |
| 49 | + |
| 50 | +def relative_correlations(lineages: list) -> dict[str, dict[str, float]]: |
| 51 | + """Model-free check on inheritance: correlation of the CDK2 activation rate between |
| 52 | + mother-daughter, sister, and first-cousin pairs born into the condition. |
| 53 | +
|
| 54 | + Under a hidden state with memory, sisters (one division apart from a shared mother) |
| 55 | + correlate more than cousins (two apart from a shared grandmother). |
| 56 | + """ |
| 57 | + pairs: dict[str, list] = {"mother_daughter": [], "sisters": [], "cousins": []} |
| 58 | + for lin in lineages: |
| 59 | + x = lin.obs[:, 0] |
| 60 | + parents, daughters = lin.edges |
| 61 | + mother = np.full(len(lin), -1) |
| 62 | + mother[daughters] = parents |
| 63 | + kids: dict[int, list[int]] = {} |
| 64 | + for p, d in zip(parents, daughters, strict=True): |
| 65 | + kids.setdefault(int(p), []).append(int(d)) |
| 66 | + if p > 0: # the root's G1 was before the drug |
| 67 | + pairs["mother_daughter"].append((x[p], x[d])) |
| 68 | + for sibs in kids.values(): |
| 69 | + if len(sibs) == 2: |
| 70 | + pairs["sisters"].append((x[sibs[0]], x[sibs[1]])) |
| 71 | + a, b = (kids.get(s, []) for s in sibs) |
| 72 | + pairs["cousins"] += [(x[i], x[j]) for i in a for j in b] |
| 73 | + out = {} |
| 74 | + for key, v in pairs.items(): |
| 75 | + arr = np.array(v, dtype=float).reshape(-1, 2) |
| 76 | + arr = arr[np.all(np.isfinite(arr), axis=1)] |
| 77 | + r = float(np.corrcoef(arr.T)[0, 1]) if arr.shape[0] > 2 else np.nan |
| 78 | + out[key] = {"r": r, "n": int(arr.shape[0])} |
| 79 | + return out |
| 80 | + |
| 81 | + |
| 82 | +def fit_drug(drug: str, n_boot: int, n_workers: int, rng, t0: float) -> dict: |
| 83 | + lineages, available = condition_lineages(drug, rng) |
| 84 | + pops, full = lineages["truncate"], lineages["keep"] |
| 85 | + res: dict = { |
| 86 | + "lineages_available": available, |
| 87 | + "cells": [int(sum(len(lin) for lin in pop)) for pop in pops], |
| 88 | + "lineages": [len(pop) for pop in pops], |
| 89 | + "depth_counts": [np.bincount(generation_depths(pop)).tolist() for pop in pops], |
| 90 | + "relatives": [relative_correlations(pop) for pop in pops], |
| 91 | + } |
| 92 | + sweep = dose_sweep(pops, list(PULSES_H), num_states=2, n_starts=8, n_jobs=n_workers, rng=rng) |
| 93 | + log(f"{drug}: dose sweep done", t0) |
| 94 | + res["LL"] = {k: float(v) for k, v in sweep.LL.items()} |
| 95 | + res["lrt"] = sweep.lrt |
| 96 | + res["emissions"] = [e.params.tolist() for e in sweep.per_dose[0].estimate.E] |
| 97 | + res["mean_lifetime_h"] = [ |
| 98 | + e.mean_lifetime() for e in sweep.per_dose[0].estimate.E if isinstance(e, StateDistribution) |
| 99 | + ] |
| 100 | + res["pi"] = [tO.estimate.pi.tolist() for tO in sweep.per_dose] |
| 101 | + res["independent_T"] = [tO.estimate.T.tolist() for tO in sweep.independent] |
| 102 | + res["point"] = summarize_T(sweep.T) |
| 103 | + |
| 104 | + res["divided_by_20h"] = division_calibration(full, sweep) |
| 105 | + res["bootstrap"], res["n_boot"] = bootstrap_summary(sweep, pops, n_boot, n_workers, rng, res["point"]) |
| 106 | + log(f"{drug}: bootstrap done", t0) |
| 107 | + |
| 108 | + # The same sweep with the roots' lifetimes taken at face value or dropped. |
| 109 | + res["root_modes"] = {} |
| 110 | + for mode in ROOT_MODES[1:]: |
| 111 | + alt = dose_sweep( |
| 112 | + lineages[mode], |
| 113 | + list(PULSES_H), |
| 114 | + num_states=2, |
| 115 | + n_starts=4, |
| 116 | + n_jobs=n_workers, |
| 117 | + rng=rng, |
| 118 | + ) |
| 119 | + res["root_modes"][mode] = { |
| 120 | + "lrt": alt.lrt, |
| 121 | + "emissions": [e.params.tolist() for e in alt.per_dose[0].estimate.E], |
| 122 | + "point": summarize_T(alt.T), |
| 123 | + "divided_by_20h": division_calibration(full, alt), |
| 124 | + } |
| 125 | + log(f"{drug}: roots={mode} done", t0) |
| 126 | + return res |
| 127 | + |
| 128 | + |
| 129 | +def run(n_boot: int = 96, n_workers: int = 44, seed: int = 0) -> dict: |
| 130 | + """Fit the drugs concurrently, each with its own share of the worker processes (the |
| 131 | + fits themselves run in those processes, so threads are enough to overlap them).""" |
| 132 | + t0 = time.time() |
| 133 | + rngs = np.random.default_rng(seed).spawn(len(DRUGS)) |
| 134 | + with ThreadPoolExecutor(len(DRUGS)) as exe: |
| 135 | + futs = { |
| 136 | + d: exe.submit(fit_drug, d, n_boot, n_workers // len(DRUGS), r, t0) for d, r in zip(DRUGS, rngs, strict=True) |
| 137 | + } |
| 138 | + out: dict = {"drugs": list(DRUGS), "pulses_h": [float(p) for p in PULSES_H]} |
| 139 | + out |= {d: f.result() for d, f in futs.items()} |
| 140 | + out["runtime_s"] = time.time() - t0 |
| 141 | + return out |
| 142 | + |
| 143 | + |
| 144 | +if __name__ == "__main__": |
| 145 | + out = run(n_boot=int(sys.argv[1]) if len(sys.argv) > 1 else 96) |
| 146 | + os.makedirs(os.path.dirname(OUTPUT), exist_ok=True) |
| 147 | + with open(OUTPUT, "w") as f: |
| 148 | + json.dump(out, f, indent=1, default=float) |
| 149 | + print(json.dumps({d: {k: out[d][k] for k in ("lrt", "point")} for d in DRUGS}, indent=1, default=float)) |
0 commit comments