{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e001",
   "metadata": {},
   "source": [
    "# Fitting with different experimental setups"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e002",
   "metadata": {},
   "source": [
    "Stable-isotope labeling experiments differ in *how* the label is introduced and removed. The `fit()` and `plot()` methods support three experimental setups via the `experiment` parameter:\n",
    "\n",
    "| Mode | Description                                            | Labeling curve `f(t)` |\n",
    "|------|--------------------------------------------------------|----------------------|\n",
    "| `WASHOUT` | Label on for a very long period, removed at $t=0$      | $W(t) = \\mathbf{s}^\\top e^{\\mathbf{M}t} \\mathbf{1}$, decays from 1 to 0 |\n",
    "| `STEP` | Label on for a defined period ($T$), removed at $t=0$  | Step-up during $[-T, 0)$, step-down for $t \\ge 0$ |\n",
    "| `PULSE` | Infinitesimally short label pulse at $t=0$            | $-\\dot{W}(t)$, the impulse response |\n",
    "\n",
    "All three setups are built from the same wash-out kernel $W(t)$, so the same compartmental model can be fitted to data from any of these designs.\n",
    "\n",
    "This notebook demonstrates:\n",
    "1. Generating synthetic labeling data for each experimental setup using a known two-state model.\n",
    "2. Fitting the model to each dataset and recovering the true parameters.\n",
    "3. Plotting the data together with the fitted labeling curves.\n",
    "4. Visualising the fitted model as a transition-rate graph with `draw_transition_matrix()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e003",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-14T09:09:52.839721229Z",
     "start_time": "2026-06-14T09:09:52.236061821Z"
    }
   },
   "outputs": [],
   "source": [
    "from symbolic_compartmental_model import SymbolicCompartmentalModel, ExperimentSetup\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "rng = np.random.default_rng(seed=42)\n",
    "\n",
    "# ── ground-truth parameters ──────────────────────────────────────────────────\n",
    "K1_TRUE = 0.5   # self-turnover of state 0\n",
    "K2_TRUE = 2.0   # self-turnover of state 1; also the contributed turnover from state 0 into state 1\n",
    "A_TRUE  = 0.3   # fraction of observable signal contributed by state 0 (slow pool)\n",
    "NOISE   = 0.03  # measurement noise standard deviation\n",
    "\n",
    "# ── model template: cascade (state 0 → state 1) with three free parameters ──\n",
    "cm = SymbolicCompartmentalModel(n_states=2)\n",
    "k1 = cm.add_parameter(\"k1\", lb=0.05, ub=5.0)\n",
    "k2 = cm.add_parameter(\"k2\", lb=0.05, ub=5.0)\n",
    "a  = cm.add_parameter(\"a\",  lb=0.01, ub=0.99)\n",
    "\n",
    "cm.contributed_turnovers = [[-k1, 0.0], [k2, -k2]]\n",
    "cm.observed_pool_weights  = [a, 1.0 - a]\n",
    "\n",
    "true_params = [K1_TRUE, K2_TRUE, A_TRUE]\n",
    "print(\"Contributed-turnover matrix M:\")\n",
    "display(cm.contributed_turnovers)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e004",
   "metadata": {},
   "source": [
    "Each entry $M_{ij}$ of the contributed-turnover matrix is defined as the **flux** between states divided by the **pool size of the receiving state**:\n",
    "\n",
    "$$M_{ij} = \\frac{\\text{flux}(j \\to i)}{\\text{pool size of state } i}, \\quad i \\ne j$$\n",
    "\n",
    "Diagonal entries $M_{ii} \\le 0$ are the self-turnover: total outflux from state $i$ per unit pool size of state $i$.\n",
    "\n",
    "For the cascade matrix above:\n",
    "- $M_{00} = -k_1$: self-turnover of state 0.\n",
    "- $M_{10} = k_2$: the flux flowing from state 0 into state 1, divided by the pool size of state 1. This equals $k_2$ (the self-turnover of state 1), so the row sum of row 1 is $k_2 - k_2 = 0$: **state 1 is internally mass-balanced** — every molecule that clears from state 1 is replaced by one coming from state 0.\n",
    "- $M_{11} = -k_2$: self-turnover of state 1.\n",
    "\n",
    "The observed signal is a mixture of both pools: $\\mathbf{s} = (a,\\, 1-a)$.\n",
    "\n",
    "We generate synthetic data by evaluating the exact labeling predictor at a set of time points and adding Gaussian noise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e005",
   "metadata": {},
   "source": [
    "## Washout experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e006",
   "metadata": {},
   "source": [
    "In a washout (chase) experiment the molecule of interest is fully labeled at $t=0$, after which labeling stops. The fraction of label remaining in the observable pool is:\n",
    "$$f(t) = W(t) = \\mathbf{s}^\\top e^{\\mathbf{M}t} \\mathbf{1},\\quad t \\ge 0$$\n",
    "This is the default mode for both `fit()` and `plot()` (`experiment=\"washout\"`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e007",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-14T09:09:53.455519537Z",
     "start_time": "2026-06-14T09:09:52.849584814Z"
    }
   },
   "outputs": [],
   "source": [
    "# ── generate synthetic washout data ─────────────────────────────────────────\n",
    "# include early time points to capture the fast component (timescale ≈ 1/k2)\n",
    "t_washout = np.array([0.2, 0.5, 1.0, 2.0, 3.0, 5.0, 8.0])\n",
    "\n",
    "washout_fn = cm._labeling_predictor(\n",
    "    x=true_params,\n",
    "    experiment=ExperimentSetup.WASHOUT,\n",
    "    period=None,\n",
    "    normalize=False,\n",
    "    use_precompiled=False,\n",
    ")\n",
    "y_washout = np.array([washout_fn(t) for t in t_washout])\n",
    "y_washout += rng.normal(0, NOISE, size=len(y_washout))\n",
    "y_washout = np.clip(y_washout, 0.0, 1.0)\n",
    "\n",
    "# ── fit ─────────────────────────────────────────────────────────────────────\n",
    "fit_washout = cm.fit(\n",
    "    t_washout, y_washout,\n",
    "    experiment=ExperimentSetup.WASHOUT,\n",
    ")\n",
    "print(fit_washout)\n",
    "\n",
    "# ── plot ─────────────────────────────────────────────────────────────────────\n",
    "fig, ax = plt.subplots(figsize=(6, 3.5), dpi=120)\n",
    "ax.scatter(t_washout, y_washout, zorder=3, label=\"synthetic data\")\n",
    "fit_washout.cm.plot(\n",
    "    \"f\", ax=ax,\n",
    "    experiment=ExperimentSetup.WASHOUT,\n",
    "    label=(\n",
    "        f\"best fit  \"\n",
    "        f\"$k_1$={fit_washout.popt[0]:.2f}, \"\n",
    "        f\"$k_2$={fit_washout.popt[1]:.2f}, \"\n",
    "        f\"$a$={fit_washout.popt[2]:.2f}\"\n",
    "    ),\n",
    ")\n",
    "ax.axhline(0, color=\"black\", linewidth=0.5)\n",
    "ax.set_title(\"Washout experiment\")\n",
    "ax.set_xlabel(\"time after label removal\")\n",
    "ax.set_ylabel(\"labeling fraction $f(t)$\")\n",
    "ax.legend()\n",
    "ax.grid(True, alpha=0.3)\n",
    "fig.tight_layout()\n",
    "if Path(\"../results\").exists():\n",
    "    fig.savefig(\"../results/example_experiment_washout.svg\")\n",
    "display(fig)\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e008",
   "metadata": {},
   "source": [
    "## Step experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e009",
   "metadata": {},
   "source": [
    "In a step experiment the label is applied for a finite period $T$ and then removed. Time is measured **from the moment labeling stops** ($t = 0$), so measurements taken *during* the labeling period have negative $t$ (as far back as $t = -T$):\n",
    "\n",
    "$$f(t) = \\begin{cases}\n",
    "0                     & t \\le -T\\quad\\text{(before labeling starts)}\\\\\n",
    "1 - W(t + T)          & -T < t < 0\\quad\\text{(step-up)}\\\\\n",
    "W(t) - W(t + T)       & t \\ge 0\\quad\\text{(wash-out / step-down)}\n",
    "\\end{cases}$$\n",
    "\n",
    "The `period` argument (in the same time units as the data) tells the model how long labeling lasted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e010",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-14T09:09:53.992963620Z",
     "start_time": "2026-06-14T09:09:53.509378478Z"
    }
   },
   "outputs": [],
   "source": [
    "T_STEP = 4.0   # labeling period\n",
    "\n",
    "# ── generate synthetic step data ──────────────────────────────────────────────\n",
    "# include early post-washout points to capture the fast component\n",
    "t_step = np.array([-3.0, -2.0, -1.0, 0.0, 0.2, 0.5, 1.0, 2.0, 3.0, 5.0])\n",
    "\n",
    "step_fn = cm._labeling_predictor(\n",
    "    x=true_params,\n",
    "    experiment=ExperimentSetup.STEP,\n",
    "    period=T_STEP,\n",
    "    normalize=False,\n",
    "    use_precompiled=False,\n",
    ")\n",
    "y_step = np.array([step_fn(t) for t in t_step])\n",
    "y_step += rng.normal(0, NOISE, size=len(y_step))\n",
    "y_step = np.clip(y_step, 0.0, 1.0)\n",
    "\n",
    "# ── fit ──────────────────────────────────────────────────────────────────────\n",
    "fit_step = cm.fit(\n",
    "    t_step, y_step,\n",
    "    experiment=ExperimentSetup.STEP,\n",
    "    period=T_STEP,\n",
    "    normalize=False,\n",
    ")\n",
    "print(fit_step)\n",
    "\n",
    "# ── plot ──────────────────────────────────────────────────────────────────────\n",
    "fig, ax = plt.subplots(figsize=(6, 3.5), dpi=120)\n",
    "ax.scatter(t_step, y_step, zorder=3, label=\"synthetic data\")\n",
    "fit_step.cm.plot(\n",
    "    \"f\", ax=ax,\n",
    "    experiment=ExperimentSetup.STEP,\n",
    "    period=T_STEP,\n",
    "    normalize=False,\n",
    "    label=(\n",
    "        f\"best fit  \"\n",
    "        f\"$k_1$={fit_step.popt[0]:.2f}, \"\n",
    "        f\"$k_2$={fit_step.popt[1]:.2f}, \"\n",
    "        f\"$a$={fit_step.popt[2]:.2f}\"\n",
    "    ),\n",
    ")\n",
    "ax.axvline(0, color=\"gray\", linestyle=\"--\", linewidth=1, label=\"label off ($t=0$)\")\n",
    "ax.axvline(-T_STEP, color=\"gray\", linestyle=\":\", linewidth=1, label=f\"label on ($t=-T$)\")\n",
    "ax.axhline(0, color=\"black\", linewidth=0.5)\n",
    "ax.set_title(f\"Step experiment ($T = {T_STEP}$)\")\n",
    "ax.set_xlabel(\"time relative to label removal\")\n",
    "ax.set_ylabel(\"labeling fraction $f(t)$\")\n",
    "ax.legend()\n",
    "ax.grid(True, alpha=0.3)\n",
    "fig.tight_layout()\n",
    "if Path(\"../results\").exists():\n",
    "    fig.savefig(\"../results/example_experiment_step.svg\")\n",
    "display(fig)\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e011",
   "metadata": {},
   "source": [
    "## Pulse experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e012",
   "metadata": {},
   "source": [
    "In a pulse experiment an infinitesimally short bolus of label is introduced at $t = 0$. The detected signal is the impulse response of the system — equivalently, the **age distribution** of molecules in the observable pool:\n",
    "$$f(t) = -\\dot{W}(t) = -\\mathbf{s}^\\top \\mathbf{M} e^{\\mathbf{M}t} \\mathbf{1},\\quad t \\ge 0$$\n",
    "\n",
    "Because $f(t)$ is a probability density, its values can exceed 1 — particularly for cascade models where the label must first pass through upstream compartments before reaching the observable pool. The normalization `normalize=True` scales the curve so that $f(0) = 1$ rather than reporting the absolute impulse amplitude. The same flag must be set consistently in both `fit()` and `plot()`.\n",
    "\n",
    "> **Physical interpretation of values > 1**: in the cascade below, newly synthesized molecules enter state 0 first. The observable signal grows for a while as molecules \"flow through\" into state 1, before eventually decaying. The peak value reflects how strongly the cascade amplifies the signal at intermediate ages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e013",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-14T09:09:56.061597151Z",
     "start_time": "2026-06-14T09:09:54.003623389Z"
    }
   },
   "outputs": [],
   "source": [
    "NOISE_PULSE = 0.04   # slightly larger noise for the pulse response\n",
    "\n",
    "# ── generate synthetic pulse data ────────────────────────────────────────────\n",
    "t_pulse = np.array([0.1, 0.3, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0])\n",
    "\n",
    "pulse_fn = cm._labeling_predictor(\n",
    "    x=true_params,\n",
    "    experiment=ExperimentSetup.PULSE,\n",
    "    period=None,\n",
    "    normalize=True,\n",
    "    use_precompiled=False,\n",
    ")\n",
    "y_pulse = np.array([pulse_fn(t) for t in t_pulse])\n",
    "y_pulse += rng.normal(0, NOISE_PULSE, size=len(y_pulse))\n",
    "# normalized pulse response is a probability density: values >= 0 but can exceed 1\n",
    "y_pulse = np.clip(y_pulse, 0.0, None)\n",
    "\n",
    "# ── fit ──────────────────────────────────────────────────────────────────────\n",
    "fit_pulse = cm.fit(\n",
    "    t_pulse, y_pulse,\n",
    "    experiment=ExperimentSetup.PULSE,\n",
    "    normalize=True,\n",
    ")\n",
    "print(fit_pulse)\n",
    "\n",
    "# ── plot ──────────────────────────────────────────────────────────────────────\n",
    "fig, ax = plt.subplots(figsize=(6, 3.5), dpi=120)\n",
    "ax.scatter(t_pulse, y_pulse, zorder=3, label=\"synthetic data\")\n",
    "fit_pulse.cm.plot(\n",
    "    \"f\", ax=ax,\n",
    "    experiment=ExperimentSetup.PULSE,\n",
    "    normalize=True,\n",
    "    label=(\n",
    "        f\"best fit  \"\n",
    "        f\"$k_1$={fit_pulse.popt[0]:.2f}, \"\n",
    "        f\"$k_2$={fit_pulse.popt[1]:.2f}, \"\n",
    "        f\"$a$={fit_pulse.popt[2]:.2f}\"\n",
    "    ),\n",
    ")\n",
    "ax.axhline(0, color=\"black\", linewidth=0.5)\n",
    "ax.set_title(\"Pulse experiment (normalized)\")\n",
    "ax.set_xlabel(\"time after pulse\")\n",
    "ax.set_ylabel(\"age distribution (normalized)\")\n",
    "ax.legend()\n",
    "ax.grid(True, alpha=0.3)\n",
    "fig.tight_layout()\n",
    "if Path(\"../results\").exists():\n",
    "    fig.savefig(\"../results/example_experiment_pulse.svg\")\n",
    "display(fig)\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e014",
   "metadata": {},
   "source": [
    "## Comparing the fitted models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e015",
   "metadata": {},
   "source": [
    "All three fits should recover parameters close to the ground truth ($k_1 = 0.5$, $k_2 = 2.0$, $a = 0.3$). The recovery quality depends on the number of data points, the noise level, and the identifiability of each experimental design.\n",
    "\n",
    "We can also inspect the fitted model graphically using `draw_transition_matrix()`, which renders the compartment network as a directed graph. Each **state node** ($s_i$) is connected to **external-input nodes** ($i_j$) when there is net external influx, and between-state edges carry the fitted transfer rates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e016",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-06-14T09:20:58.623150907Z",
     "start_time": "2026-06-14T09:20:58.233404885Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 2, figsize=(8, 8), dpi=120)\n",
    "\n",
    "for ax, params, title in zip(\n",
    "    axs.flat,\n",
    "    [true_params, fit_washout.popt, fit_step.popt, fit_pulse.popt],\n",
    "    [\"True\", \"Washout\", \"Step\", \"Pulse\"],\n",
    "):\n",
    "    cm.draw_transition_matrix(ax=ax, x=params)\n",
    "    k1_fit, k2_fit, a_fit = params\n",
    "    ax.set_title(\n",
    "        f\"{title}\\n $k_1$={k1_fit:.2f}, $k_2$={k2_fit:.2f}, $a$={a_fit:.2f}\",\n",
    "        fontsize=8\n",
    "    )\n",
    "    ax.set_axis_off()\n",
    "\n",
    "fig.tight_layout()\n",
    "if Path(\"../results\").exists():\n",
    "    fig.savefig(\"../results/example_experiment_transition_matrix.svg\")\n",
    "display(fig)\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e017",
   "metadata": {},
   "source": [
    "Each graph shows the same two-edge cascade topology. Edge labels are **contributed turnovers** (flux divided by receiving pool size):\n",
    "\n",
    "- $i_0 \\to s_0$  unlabeled external input to state 0 (value = $k_1$), replenishing state 0's clearance.\n",
    "- $s_0 \\to s_1$  contributed turnover from state 0 into state 1 ($M_{10} = k_2$) — the flux from state 0 into state 1, divided by the pool size of state 1.\n",
    "\n",
    "Because state 1 is internally mass-balanced (row 1 sums to zero), there is no external-input node for state 1."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbformat_minor": 5
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}