adesomics.population.population_model

  1import numpy as np
  2import pandas as pd
  3
  4from adesomics.io.metadata import Metadata
  5from adesomics.plotting.style import annotate_study, group_colors, short_labels
  6import matplotlib.pyplot as plt
  7from adesomics.plotting.style import group_colors, annotate_study, GREY, INK
  8
  9from pathlib import Path
 10
 11def _save(fig, filename: str, outdir: str = "results") -> Path:
 12    out = Path(outdir)
 13    out.mkdir(parents=True, exist_ok=True)      
 14    path = (out / filename).resolve()           
 15    fig.savefig(path, dpi=200, bbox_inches="tight")
 16    print(f"Saved plot to {path}")              
 17    return path
 18
 19class Population:
 20    """Exponential growth of one sample, driven by its FBA growth rate.
 21
 22    Parameters
 23    ----------
 24    growth_rate
 25        Specific growth rate mu (1/h). If omitted, it is taken from the
 26        model's FBA objective on first access.
 27    """
 28
 29    def __init__(self, model=None, n0: float = 1.0,
 30                 growth_rate: float | None = None, label: str | None = None):
 31        if model is None and growth_rate is None:
 32            raise ValueError("Need either a model or an explicit growth_rate")
 33        self.model = model
 34        self.n0 = n0
 35        self._growth_rate = growth_rate
 36        self.label = label
 37
 38    @property
 39    def growth_rate(self) -> float:
 40        if self._growth_rate is None:
 41            self._growth_rate = self.model.optimize().objective_value
 42        return self._growth_rate
 43
 44    @property
 45    def doubling_time(self) -> float:
 46        """Doubling time in hours; inf for a non-growing population."""
 47        mu = self.growth_rate
 48        return np.log(2) / mu if mu > 0 else np.inf
 49
 50    def simulate(self, t_max: float, n_points: int = 100):
 51        t = np.linspace(0, t_max, n_points)
 52        return t, self.n0 * np.exp(self.growth_rate * t)
 53
 54    def plot_growth_curve(self, t_max: float, n_points: int = 100,
 55                          ax=None, **kwargs):
 56        import matplotlib.pyplot as plt
 57
 58        t, n = self.simulate(t_max, n_points)
 59        ax = ax or plt.subplots(figsize=(7, 5))[1]
 60        kwargs.setdefault("label", self.label or f"mu = {self.growth_rate:.3f} 1/h")
 61        ax.plot(t, n, **kwargs)
 62        ax.set(xlabel="Time (h)", ylabel="Population size")
 63        return ax
 64
 65    def __repr__(self):
 66        return f"<Population {self.label!r}: mu={self.growth_rate:.3f} 1/h>"
 67
 68
 69class PopulationSet:
 70    """One Population per sample, tied to the experiment's Metadata.
 71
 72    Colours and labels come from adesomics.plotting.style, so growth curves
 73    match the DE plots of the same experiment.
 74    """
 75
 76    def __init__(self, populations: dict[str, Population], metadata: Metadata):
 77        if metadata is not None:
 78            missing = set(populations) - set(metadata.samples)
 79            if missing:
 80                raise ValueError(f"Not in metadata: {sorted(missing)}")
 81
 82            self.metadata = metadata.subset_samples(list(populations))
 83            self.populations = {s: populations[s] for s in self.metadata.samples}
 84        else:
 85            print("Populationmodel currently only works with linked metadata as input.")
 86            raise SystemExit(1) 
 87            
 88    # --- construction ----------------------------------------------------
 89
 90    @classmethod
 91    def from_connector(cls, connector, metadata: Metadata,
 92                       samples: list[str] | None = None, n0: float = 1.0):
 93        """Build one Population per sample from E-Flux-constrained solutions."""
 94        samples = samples if samples is not None else list(connector.tpm.columns)
 95        solutions = connector.compare(samples)
 96
 97        pops = {
 98            s: Population(n0=n0,
 99                          growth_rate=solutions[s].objective_value,
100                          label=s)
101            for s in samples
102        }
103        return cls(pops, metadata)
104
105    # --- access ----------------------------------------------------------
106
107    @property
108    def growth_rates(self) -> pd.Series:
109        return pd.Series({s: p.growth_rate for s, p in self.populations.items()},
110                         name="mu")
111
112    def summary(self) -> pd.DataFrame:
113        """Growth rate and doubling time per sample, with its condition."""
114        df = pd.DataFrame({
115            "condition": self.metadata.groups,
116            "mu": self.growth_rates,
117            "doubling_time_h": {s: p.doubling_time
118                                for s, p in self.populations.items()},
119        })
120        return df.loc[self.metadata.samples]
121
122    def by_group(self) -> pd.DataFrame:
123        """Mean and spread of mu per condition (replicates collapsed)."""
124        return (self.summary()
125                .groupby("condition")["mu"]
126                .agg(["mean", "std", "count"]))
127
128    # --- plotting --------------------------------------------------------
129
130    def plot_growth_curves(self, t_max: float, n_points: int = 100,
131                           ax=None, by_group: bool = True):
132        ax = ax or plt.subplots(figsize=(7, 5))[1]
133        lut = group_colors(self.metadata)
134        seen = set()
135
136        for sample in self.metadata.samples:
137            group = self.metadata.groups[sample]
138            if by_group:
139                label = group if group not in seen else None
140                seen.add(group)
141            else:
142                label = sample
143
144            self.populations[sample].plot_growth_curve(
145                t_max, n_points, ax=ax, color=lut[group], label=label,
146            )
147
148        ax.set(title="FBA-predicted exponential growth")
149        ax.legend(frameon=False, fontsize=7)
150        annotate_study(ax, self.metadata)
151        _save(ax.figure, "growth_curves.png", "results")
152        return ax
153
154    def plot_growth_rates(self, ax=None, save=None):
155        summ = self.summary()
156        by = self.by_group()
157        lut = group_colors(self.metadata)
158        order = list(by.index)
159
160        ax = ax or plt.subplots(figsize=(8, 5))[1]
161        label_x = summ["mu"].max() * 1.05
162        for i, group in enumerate(order):
163            members = summ.index[summ["condition"] == group]
164            ax.barh(i, by.loc[group, "mean"], color=lut[group], alpha=0.85,
165                    height=0.62, zorder=2)
166            ax.scatter(summ.loc[members, "mu"], [i] * len(members),
167                    color=INK, s=30, zorder=3)
168            td = np.log(2) / by.loc[group, "mean"]
169            ax.text(label_x, i, f"t$_d$ = {td:.1f} h",
170                va="center", fontsize=10, color=GREY)
171
172        ax.set_yticks(range(len(order)))
173        ax.set_yticklabels([g.replace("_", " ") for g in order])
174        ax.set(xlabel="specific growth rate  $\\mu$  (1/h)",
175            title="Growth rate by condition")
176        ax.margins(x=0.22)
177        ax.grid(axis="y", visible=False)
178        annotate_study(ax, self.metadata)
179        _save(ax.figure, "Bar_GrowthRates.png", "results")
180        return ax
181
182    def __len__(self):
183        return len(self.populations)
184
185    def __repr__(self):
186        return (f"<PopulationSet: {len(self)} samples, "
187                f"{self.metadata.groups.nunique()} conditions>")