{
"cells": [
{
"cell_type": "markdown",
"id": "1c648a7c",
"metadata": {},
"source": [
"# Conditional Diffusion Method\n",
"\n",
"This notebook implements a thesis-specific **conditional diffusion model** for single-gene perturbation prediction on single-cell RNA-seq data.\n",
"\n",
"Core idea:\n",
"\n",
"1. Encode control and perturbed cells into a latent space.\n",
"2. Build pseudo paired endpoints with mini-batch optimal transport.\n",
"3. Train a perturbation-conditioned denoiser to reverse a diffusion process on perturbed latent states.\n",
"4. Decode the sampled latent state back to expression space.\n",
"\n",
"This gives a distributional perturbation model parallel to the flow-matching notebook, but with a diffusion-based generative mechanism.\n"
]
},
{
"cell_type": "markdown",
"id": "54e1b660",
"metadata": {},
"source": [
"## Method Summary\n",
"\n",
"Let `x_ctrl` be a control-cell expression vector and let `p` be a single-gene perturbation.\n",
"\n",
"The model uses:\n",
"\n",
"- an encoder `E` to map cells into a latent state `z`\n",
"- a condition encoder for perturbation gene and batch context\n",
"- a denoiser `eps_theta(z_t, t, z_ctrl, h)` trained in DDPM style\n",
"- a decoder `D` to map generated latent states back to expression space\n",
"\n",
"Training target:\n",
"\n",
"- learn to denoise noisy perturbed latent states conditioned on the matched control latent state\n",
"\n",
"Inference:\n",
"\n",
"- start from Gaussian noise\n",
"- reverse diffuse to a perturbation latent state conditioned on the control cell and perturbation\n",
"- decode to gene expression\n"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "fbe26970",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Project root: /root/final\n",
"Input h5ad: /root/final/data/processed/scPerturb/prediction_ready/replogle_k562_essential_single_gene_prediction_ready_full.h5ad\n",
"Output dir: /root/final/baseline/outputs/conditional_diffusion_method\n",
"Device: cuda\n"
]
}
],
"source": [
"from pathlib import Path\n",
"from copy import deepcopy\n",
"import json\n",
"import math\n",
"import random\n",
"\n",
"import anndata as ad\n",
"import matplotlib as mpl\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import pandas as pd\n",
"import scipy.sparse as sp\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"\n",
"\n",
"def find_project_root(start: Path | None = None) -> Path:\n",
" start = (start or Path.cwd()).resolve()\n",
" for path in [start, *start.parents]:\n",
" if (path / \"src\" / \"data\").exists() and (path / \"baseline\").exists():\n",
" return path\n",
" return start\n",
"\n",
"\n",
"PROJECT_ROOT = find_project_root()\n",
"DATASET_NAME = \"ReplogleWeissman2022_K562_essential\"\n",
"METHOD_NAME = \"conditional_diffusion_method\"\n",
"RANDOM_SEED = 42\n",
"\n",
"SPLIT_STRATEGY = \"seen_cell_split\" # choices: \"seen_cell_split\", \"use_existing\"\n",
"TRAIN_FRAC = 0.80\n",
"VAL_FRAC = 0.10\n",
"EVAL_SPLIT = \"test\"\n",
"CONTROL_SPLIT = \"train\"\n",
"TRAIN_SPLIT = \"train\"\n",
"VAL_SPLIT = \"val\"\n",
"\n",
"USE_FEATURE_GENES = True\n",
"FEATURE_GENE_COLUMNS = [\"is_feature_gene\", \"highly_variable\"]\n",
"\n",
"MAX_TRAIN_CONDITIONS = None\n",
"MAX_EVAL_GENES = None\n",
"MAX_EVAL_CELLS_PER_GENE = None\n",
"\n",
"LATENT_DIM = 64\n",
"HIDDEN_DIM = 256\n",
"CONDITION_DIM = 128\n",
"TIME_DIM = 64\n",
"DROPOUT = 0.10\n",
"NUM_DENOISER_BLOCKS = 4\n",
"\n",
"BATCH_SIZE = 128\n",
"EPOCHS = 20\n",
"STEPS_PER_EPOCH = 200\n",
"VALIDATION_STEPS = 32\n",
"LEARNING_RATE = 1e-3\n",
"WEIGHT_DECAY = 1e-5\n",
"GRAD_CLIP_NORM = 1.0\n",
"\n",
"RECON_WEIGHT = 1.0\n",
"DIFFUSION_WEIGHT = 1.0\n",
"ENDPOINT_WEIGHT = 0.5\n",
"MMD_WEIGHT = 0.1\n",
"MEAN_WEIGHT = 1.0\n",
"\n",
"DIFFUSION_STEPS = 50\n",
"BETA_START = 1e-4\n",
"BETA_END = 2e-2\n",
"SINKHORN_EPSILON = 0.5\n",
"SINKHORN_ITERS = 30\n",
"PRED_BATCH_SIZE = 256\n",
"\n",
"INPUT_CANDIDATES = [\n",
" PROJECT_ROOT / \"data\" / \"processed\" / \"scPerturb\" / \"prediction_ready\" / \"replogle_k562_essential_single_gene_prediction_ready_full.h5ad\",\n",
" PROJECT_ROOT / \"data\" / \"processed\" / \"scPerturb\" / \"prediction_ready\" / \"replogle_k562_essential_single_gene_prediction_ready.h5ad\",\n",
" PROJECT_ROOT / \"data\" / \"processed\" / \"scPerturb\" / \"prediction_ready\" / \"replogle_k562_essential_single_gene_prediction_ready_fast_debug.h5ad\",\n",
"]\n",
"INPUT_PATH = next((path for path in INPUT_CANDIDATES if path.exists()), None)\n",
"if INPUT_PATH is None:\n",
" raise FileNotFoundError(\"Could not find a prediction-ready Replogle h5ad. Checked:\\n\" + \"\\n\".join(map(str, INPUT_CANDIDATES)))\n",
"\n",
"OUTPUT_DIR = PROJECT_ROOT / \"baseline\" / \"outputs\" / METHOD_NAME\n",
"OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
"MODEL_PATH = OUTPUT_DIR / \"conditional_diffusion_model.pt\"\n",
"\n",
"DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
"\n",
"np.random.seed(RANDOM_SEED)\n",
"random.seed(RANDOM_SEED)\n",
"torch.manual_seed(RANDOM_SEED)\n",
"if torch.cuda.is_available():\n",
" torch.cuda.manual_seed_all(RANDOM_SEED)\n",
"if torch.backends.cudnn.is_available():\n",
" torch.backends.cudnn.deterministic = True\n",
" torch.backends.cudnn.benchmark = False\n",
"\n",
"print(f\"Project root: {PROJECT_ROOT}\")\n",
"print(f\"Input h5ad: {INPUT_PATH}\")\n",
"print(f\"Output dir: {OUTPUT_DIR}\")\n",
"print(f\"Device: {DEVICE}\")\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "diffusion-science-plots",
"metadata": {},
"outputs": [],
"source": [
"SCIENCE_COLORS = {\n",
" \"blue\": \"#0072B2\",\n",
" \"orange\": \"#E69F00\",\n",
" \"green\": \"#009E73\",\n",
" \"magenta\": \"#CC79A7\",\n",
" \"dark\": \"#333333\",\n",
" \"gray\": \"#999999\",\n",
"}\n",
"\n",
"\n",
"def setup_science_style() -> None:\n",
" mpl.rcParams.update(\n",
" {\n",
" \"figure.dpi\": 110,\n",
" \"savefig.dpi\": 300,\n",
" \"font.family\": \"sans-serif\",\n",
" \"font.sans-serif\": [\"DejaVu Sans\", \"Helvetica\", \"Arial\"],\n",
" \"font.size\": 9,\n",
" \"axes.titlesize\": 10,\n",
" \"axes.labelsize\": 9,\n",
" \"axes.linewidth\": 0.6,\n",
" \"axes.edgecolor\": SCIENCE_COLORS[\"dark\"],\n",
" \"axes.labelcolor\": SCIENCE_COLORS[\"dark\"],\n",
" \"xtick.labelsize\": 8,\n",
" \"ytick.labelsize\": 8,\n",
" \"xtick.major.width\": 0.6,\n",
" \"ytick.major.width\": 0.6,\n",
" \"legend.frameon\": False,\n",
" \"legend.fontsize\": 8,\n",
" \"axes.spines.top\": False,\n",
" \"axes.spines.right\": False,\n",
" \"figure.facecolor\": \"white\",\n",
" \"axes.facecolor\": \"white\",\n",
" \"grid.alpha\": 0.3,\n",
" \"grid.linewidth\": 0.4,\n",
" \"lines.linewidth\": 1.25,\n",
" \"pdf.fonttype\": 42,\n",
" \"ps.fonttype\": 42,\n",
" }\n",
" )\n",
"\n",
"\n",
"FIGURE_DIR = OUTPUT_DIR / \"figures\"\n",
"FIGURE_DIR.mkdir(parents=True, exist_ok=True)\n",
"setup_science_style()\n",
"\n",
"\n",
"def plot_training_history(history: pd.DataFrame, tag: str) -> None:\n",
" setup_science_style()\n",
" h = history.copy()\n",
" if h.empty:\n",
" return\n",
" fig, axes = plt.subplots(1, 2, figsize=(7.0, 2.45), constrained_layout=True)\n",
" e = h[\"epoch\"].to_numpy()\n",
" ax = axes[0]\n",
" if \"train_total\" in h.columns:\n",
" ax.plot(\n",
" e,\n",
" h[\"train_total\"],\n",
" \"o-\",\n",
" ms=3,\n",
" lw=1.1,\n",
" color=SCIENCE_COLORS[\"blue\"],\n",
" label=\"Train (total)\",\n",
" )\n",
" if \"val_total\" in h.columns and h[\"val_total\"].notna().any():\n",
" ax.plot(\n",
" e,\n",
" h[\"val_total\"],\n",
" \"s-\",\n",
" ms=3,\n",
" lw=1.1,\n",
" color=SCIENCE_COLORS[\"orange\"],\n",
" label=\"Validation (total)\",\n",
" )\n",
" best_i = int(h[\"val_total\"].idxmin())\n",
" best_ep = float(h.loc[best_i, \"epoch\"])\n",
" best_v = float(h.loc[best_i, \"val_total\"])\n",
" ax.scatter(\n",
" [best_ep],\n",
" [best_v],\n",
" s=55,\n",
" zorder=5,\n",
" facecolors=\"none\",\n",
" edgecolors=SCIENCE_COLORS[\"dark\"],\n",
" linewidths=1.1,\n",
" label=\"Best val\",\n",
" )\n",
" ax.set_xlabel(\"Epoch\")\n",
" ax.set_ylabel(\"Loss\")\n",
" ax.legend(loc=\"upper right\", handlelength=2.0)\n",
"\n",
" ax2 = axes[1]\n",
" for col, lab, c in (\n",
" (\"train_diffusion\", \"Train diffusion\", SCIENCE_COLORS[\"green\"]),\n",
" (\"train_recon\", \"Train reconstruction\", SCIENCE_COLORS[\"magenta\"]),\n",
" (\"train_mmd\", \"Train MMD\", \"#6A3D9A\"),\n",
" (\"train_mean\", \"Train mean\", \"#D55E00\"),\n",
" ):\n",
" if col in h.columns:\n",
" ax2.plot(e, h[col], \"o-\", ms=2.5, lw=1.0, color=c, label=lab)\n",
" ax2.set_xlabel(\"Epoch\")\n",
" ax2.set_ylabel(\"Component loss\")\n",
" ax2.set_yscale(\"symlog\", linthresh=1e-4)\n",
" ax2.legend(loc=\"upper right\", handlelength=2.0)\n",
"\n",
" fig.suptitle(\n",
" \"Training dynamics — conditional diffusion\",\n",
" fontsize=10.5,\n",
" color=SCIENCE_COLORS[\"dark\"],\n",
" y=1.02,\n",
" )\n",
" for ext in (\"png\", \"pdf\"):\n",
" fig.savefig(FIGURE_DIR / f\"training_history_{tag}.{ext}\", bbox_inches=\"tight\", facecolor=\"white\")\n",
" plt.show()\n",
"\n",
"\n",
"def plot_evaluation_metrics(\n",
" metrics: pd.DataFrame,\n",
" tag: str,\n",
" train_genes: list[str] | None = None,\n",
") -> None:\n",
" setup_science_style()\n",
" if metrics.empty:\n",
" return\n",
" m = metrics.copy()\n",
" if train_genes is not None:\n",
" seen = set(train_genes)\n",
" m[\"cohort\"] = np.where(m[\"perturbation_gene\"].isin(seen), \"Seen at train\", \"Unseen at train\")\n",
" else:\n",
" m[\"cohort\"] = \"All\"\n",
"\n",
" fig = plt.figure(figsize=(7.2, 5.9), constrained_layout=True)\n",
" gs = fig.add_gridspec(2, 2, height_ratios=[1.05, 1.0], hspace=0.28, wspace=0.28)\n",
"\n",
" ax1 = fig.add_subplot(gs[0, 0])\n",
" if m[\"cohort\"].nunique() > 1:\n",
" palette = {\"Seen at train\": SCIENCE_COLORS[\"blue\"], \"Unseen at train\": SCIENCE_COLORS[\"orange\"]}\n",
" for lab, g in m.groupby(\"cohort\"):\n",
" s = 14 + np.clip(g[\"n_cells\"].to_numpy(dtype=float), 0, 120) * 0.12\n",
" ax1.scatter(\n",
" g[\"mse_mean\"],\n",
" g[\"pearson_mean\"],\n",
" s=s,\n",
" alpha=0.78,\n",
" edgecolors=\"none\",\n",
" c=palette.get(lab, SCIENCE_COLORS[\"gray\"]),\n",
" label=lab,\n",
" )\n",
" ax1.legend(loc=\"lower left\")\n",
" else:\n",
" ax1.scatter(\n",
" m[\"mse_mean\"],\n",
" m[\"pearson_mean\"],\n",
" s=22,\n",
" alpha=0.75,\n",
" edgecolors=\"none\",\n",
" c=SCIENCE_COLORS[\"blue\"],\n",
" )\n",
" ax1.set_xlabel(\"MSE (predicted vs true, gene means)\")\n",
" ax1.set_ylabel(\"Pearson r (gene means)\")\n",
" ax1.set_xscale(\"log\")\n",
" ax1.text(0.02, 0.98, \"a\", transform=ax1.transAxes, fontsize=11, fontweight=\"bold\", va=\"top\")\n",
"\n",
" ax2 = fig.add_subplot(gs[0, 1])\n",
" m_sorted = m.sort_values(\"pearson_mean\", ascending=True)\n",
" max_rows = 36\n",
" if len(m_sorted) > max_rows:\n",
" head_n = max_rows // 2\n",
" plot_df = pd.concat([m_sorted.head(head_n), m_sorted.tail(head_n)], axis=0)\n",
" else:\n",
" plot_df = m_sorted\n",
" y = np.arange(len(plot_df))\n",
" ax2.barh(y, plot_df[\"pearson_mean\"], height=0.68, color=SCIENCE_COLORS[\"blue\"], alpha=0.88)\n",
" ax2.set_yticks(y)\n",
" ax2.set_yticklabels(plot_df[\"perturbation_gene\"], fontsize=5.8)\n",
" ax2.set_xlabel(\"Pearson r (gene means)\")\n",
" ax2.set_xlim(0, 1.02)\n",
" ax2.text(0.02, 0.98, \"b\", transform=ax2.transAxes, fontsize=11, fontweight=\"bold\", va=\"top\")\n",
"\n",
" ax3 = fig.add_subplot(gs[1, 0])\n",
" d = m[\"delta_l2\"].dropna().to_numpy()\n",
" if d.size:\n",
" nbin = int(np.clip(len(d) // 2, 10, 32))\n",
" ax3.hist(\n",
" d,\n",
" bins=nbin,\n",
" color=SCIENCE_COLORS[\"green\"],\n",
" alpha=0.88,\n",
" edgecolor=\"white\",\n",
" linewidth=0.35,\n",
" )\n",
" ax3.set_xlabel(\"L2 distance (pred delta vs true delta)\")\n",
" ax3.set_ylabel(\"Perturbations\")\n",
" ax3.text(0.02, 0.98, \"c\", transform=ax3.transAxes, fontsize=11, fontweight=\"bold\", va=\"top\")\n",
"\n",
" ax4 = fig.add_subplot(gs[1, 1])\n",
" ax4.scatter(\n",
" m[\"pearson_mean\"],\n",
" m[\"pearson_delta\"],\n",
" s=20,\n",
" c=SCIENCE_COLORS[\"orange\"],\n",
" alpha=0.72,\n",
" edgecolors=\"none\",\n",
" )\n",
" pm = m[\"pearson_mean\"].to_numpy(dtype=float)\n",
" pd_ = m[\"pearson_delta\"].to_numpy(dtype=float)\n",
" lo = float(np.nanmin([np.nanmin(pm), np.nanmin(pd_), 0.0]))\n",
" hi = float(np.nanmax([np.nanmax(pm), np.nanmax(pd_), 1.0]))\n",
" pad = 0.03 * (hi - lo + 1e-6)\n",
" ax4.plot([lo, hi], [lo, hi], ls=\"--\", color=SCIENCE_COLORS[\"gray\"], lw=0.9, label=\"Identity\")\n",
" ax4.set_xlim(lo - pad, hi + pad)\n",
" ax4.set_ylim(lo - pad, hi + pad)\n",
" ax4.set_xlabel(\"Pearson r (means)\")\n",
" ax4.set_ylabel(\"Pearson r (delta vs control)\")\n",
" ax4.legend(loc=\"lower right\")\n",
" ax4.text(0.02, 0.98, \"d\", transform=ax4.transAxes, fontsize=11, fontweight=\"bold\", va=\"top\")\n",
"\n",
" fig.suptitle(\n",
" \"Held-out evaluation — per-perturbation metrics\",\n",
" fontsize=10.5,\n",
" color=SCIENCE_COLORS[\"dark\"],\n",
" )\n",
" for ext in (\"png\", \"pdf\"):\n",
" fig.savefig(FIGURE_DIR / f\"evaluation_metrics_{tag}.{ext}\", bbox_inches=\"tight\", facecolor=\"white\")\n",
" plt.show()\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "9c539fd8",
"metadata": {},
"outputs": [],
"source": [
"def normalize_bool(series: pd.Series) -> pd.Series:\n",
" if pd.api.types.is_bool_dtype(series):\n",
" return series.fillna(False).astype(bool)\n",
" text = series.astype(str).str.lower().str.strip()\n",
" return text.isin([\"true\", \"1\", \"yes\", \"y\"])\n",
"\n",
"\n",
"def require_columns(obs: pd.DataFrame, columns: list[str]) -> None:\n",
" missing = [col for col in columns if col not in obs.columns]\n",
" if missing:\n",
" raise KeyError(f\"Missing required obs columns: {missing}\")\n",
"\n",
"\n",
"def take_dense_rows(X, idx: np.ndarray) -> np.ndarray:\n",
" out = X[idx]\n",
" if sp.issparse(out):\n",
" out = out.toarray()\n",
" return np.asarray(out, dtype=np.float32)\n",
"\n",
"\n",
"def mean_vector(X) -> np.ndarray:\n",
" if sp.issparse(X):\n",
" return np.asarray(X.mean(axis=0)).ravel().astype(np.float32)\n",
" return np.asarray(X, dtype=np.float32).mean(axis=0)\n",
"\n",
"\n",
"def safe_pearson(x: np.ndarray, y: np.ndarray) -> float:\n",
" x = np.asarray(x).ravel()\n",
" y = np.asarray(y).ravel()\n",
" if x.size < 2 or np.std(x) == 0 or np.std(y) == 0:\n",
" return np.nan\n",
" return float(np.corrcoef(x, y)[0, 1])\n",
"\n",
"\n",
"def split_array_indices(indices: np.ndarray, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
" indices = np.asarray(indices)\n",
" if len(indices) == 0:\n",
" return indices, indices, indices\n",
" shuffled = indices.copy()\n",
" rng.shuffle(shuffled)\n",
" if len(shuffled) < 3:\n",
" return np.sort(shuffled), np.array([], dtype=indices.dtype), np.array([], dtype=indices.dtype)\n",
" n_train = max(1, int(round(len(shuffled) * TRAIN_FRAC)))\n",
" n_val = int(round(len(shuffled) * VAL_FRAC))\n",
" n_train = min(n_train, len(shuffled) - 2)\n",
" n_val = max(1, min(n_val, len(shuffled) - n_train - 1))\n",
" train = np.sort(shuffled[:n_train])\n",
" val = np.sort(shuffled[n_train : n_train + n_val])\n",
" test = np.sort(shuffled[n_train + n_val :])\n",
" if len(test) == 0:\n",
" test = val[-1:].copy()\n",
" val = val[:-1]\n",
" if len(val) == 0:\n",
" val = train[-1:].copy()\n",
" train = train[:-1]\n",
" return train, val, test\n",
"\n",
"\n",
"def assign_seen_cell_split(obs: pd.DataFrame, seed: int) -> pd.Categorical:\n",
" rng = np.random.default_rng(seed)\n",
" split = np.full(len(obs), \"train\", dtype=object)\n",
" is_control = normalize_bool(obs[\"is_control\"]).to_numpy(dtype=bool)\n",
" genes = obs[\"perturbation_gene\"].astype(str).to_numpy()\n",
"\n",
" control_idx = np.where(is_control)[0]\n",
" train_idx, val_idx, test_idx = split_array_indices(control_idx, rng)\n",
" split[train_idx] = \"train\"\n",
" split[val_idx] = \"val\"\n",
" split[test_idx] = \"test\"\n",
"\n",
" for gene in sorted(set(genes[~is_control])):\n",
" idx = np.where((~is_control) & (genes == gene))[0]\n",
" train_idx, val_idx, test_idx = split_array_indices(idx, rng)\n",
" split[train_idx] = \"train\"\n",
" split[val_idx] = \"val\"\n",
" split[test_idx] = \"test\"\n",
"\n",
" return pd.Categorical(split, categories=[\"train\", \"val\", \"test\"], ordered=True)\n",
"\n",
"\n",
"def choose_feature_mask(var: pd.DataFrame) -> np.ndarray:\n",
" if not USE_FEATURE_GENES:\n",
" return np.ones(var.shape[0], dtype=bool)\n",
" for column in FEATURE_GENE_COLUMNS:\n",
" if column in var.columns:\n",
" mask = var[column].astype(bool).to_numpy()\n",
" if mask.any():\n",
" print(f\"Using feature mask from var['{column}']: {int(mask.sum())} genes\")\n",
" return mask\n",
" print(\"Feature-gene columns not found or empty; using all genes.\")\n",
" return np.ones(var.shape[0], dtype=bool)\n",
"\n",
"\n",
"def median_heuristic_sigma(x: torch.Tensor, y: torch.Tensor, max_points: int = 256) -> float:\n",
" z = torch.cat([x[:max_points], y[:max_points]], dim=0)\n",
" if z.shape[0] < 2:\n",
" return 1.0\n",
" d2 = torch.cdist(z, z, p=2).pow(2)\n",
" mask = ~torch.eye(d2.shape[0], dtype=torch.bool, device=d2.device)\n",
" valid = d2[mask]\n",
" valid = valid[valid > 0]\n",
" if valid.numel() == 0:\n",
" return 1.0\n",
" return float(torch.sqrt(valid.median()).item())\n",
"\n",
"\n",
"def rbf_mmd_loss(x: torch.Tensor, y: torch.Tensor, sigma: float | None = None) -> torch.Tensor:\n",
" if x.shape[0] < 2 or y.shape[0] < 2:\n",
" return x.new_tensor(0.0)\n",
" sigma = sigma or median_heuristic_sigma(x, y)\n",
" gamma = 1.0 / (2.0 * sigma * sigma + 1e-8)\n",
" d_xx = torch.cdist(x, x, p=2).pow(2)\n",
" d_yy = torch.cdist(y, y, p=2).pow(2)\n",
" d_xy = torch.cdist(x, y, p=2).pow(2)\n",
" k_xx = torch.exp(-gamma * d_xx)\n",
" k_yy = torch.exp(-gamma * d_yy)\n",
" k_xy = torch.exp(-gamma * d_xy)\n",
" return k_xx.mean() + k_yy.mean() - 2.0 * k_xy.mean()\n",
"\n",
"\n",
"def sinkhorn_barycentric_targets(x0: torch.Tensor, x1: torch.Tensor, epsilon: float, n_iters: int) -> torch.Tensor:\n",
" cost = torch.cdist(x0, x1, p=2).pow(2)\n",
" kernel = torch.exp(-cost / max(epsilon, 1e-4)).clamp_min(1e-8)\n",
"\n",
" a = torch.full((x0.shape[0],), 1.0 / x0.shape[0], device=x0.device, dtype=x0.dtype)\n",
" b = torch.full((x1.shape[0],), 1.0 / x1.shape[0], device=x1.device, dtype=x1.dtype)\n",
" u = torch.ones_like(a)\n",
" v = torch.ones_like(b)\n",
"\n",
" for _ in range(n_iters):\n",
" u = a / (kernel @ v + 1e-8)\n",
" v = b / (kernel.t() @ u + 1e-8)\n",
"\n",
" plan = u[:, None] * kernel * v[None, :]\n",
" row_mass = plan.sum(dim=1, keepdim=True).clamp_min(1e-8)\n",
" return (plan / row_mass) @ x1\n",
"\n",
"\n",
"def build_aggregate_frame(metrics: pd.DataFrame) -> pd.DataFrame:\n",
" numeric = metrics.select_dtypes(include=[np.number])\n",
" if numeric.empty:\n",
" return pd.DataFrame()\n",
" return numeric.agg([\"mean\", \"median\"]).T\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "2485e6fa",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Using feature mask from var['is_feature_gene']: 4334 genes\n",
"adata: AnnData object with n_obs × n_vars = 310385 × 4334\n",
" obs: 'batch', 'gene', 'gene_id', 'transcript', 'gene_transcript', 'guide_id', 'percent_mito', 'UMI_count', 'z_gemgroup_UMI', 'core_scale_factor', 'core_adjusted_UMI_count', 'disease', 'cancer', 'cell_line', 'sex', 'age', 'perturbation', 'organism', 'perturbation_type', 'tissue_type', 'ncounts', 'ngenes', 'nperts', 'percent_ribo', 'perturbation_label', 'perturbation_gene', 'is_control', 'is_single_perturbation', 'split', 'condition', 'batch_model', 'source_split'\n",
" var: 'chr', 'start', 'end', 'class', 'strand', 'length', 'in_matrix', 'mean', 'std', 'cv', 'fano', 'ensembl_id', 'ncounts', 'ncells', 'highly_variable', 'highly_variable_rank', 'means', 'variances', 'variances_norm', 'is_target_gene', 'is_feature_gene'\n",
" uns: 'hvg', 'log1p', 'prediction_ready'\n",
" layers: 'counts'\n",
"split counts:\n",
"split\n",
"train 248311\n",
"test 31046\n",
"val 31028\n",
"Name: count, dtype: int64\n",
"train conditions: 2056\n",
"val conditions: 2056\n",
"eval conditions: 2056\n",
"eval genes missing from training: 0\n"
]
}
],
"source": [
"adata = ad.read_h5ad(INPUT_PATH)\n",
"require_columns(adata.obs, [\"perturbation_gene\", \"is_control\", \"split\"])\n",
"\n",
"adata.obs[\"is_control\"] = normalize_bool(adata.obs[\"is_control\"])\n",
"adata.obs[\"perturbation_gene\"] = adata.obs[\"perturbation_gene\"].astype(str).str.strip()\n",
"adata.obs[\"perturbation_gene\"] = adata.obs[\"perturbation_gene\"].mask(adata.obs[\"is_control\"], \"__ctrl__\")\n",
"adata.obs[\"condition\"] = adata.obs[\"perturbation_gene\"].astype(str)\n",
"adata.obs[\"batch_model\"] = adata.obs[\"batch\"].astype(str) if \"batch\" in adata.obs.columns else \"global\"\n",
"\n",
"feature_mask = choose_feature_mask(adata.var)\n",
"adata = adata[:, feature_mask].copy()\n",
"\n",
"if SPLIT_STRATEGY == \"seen_cell_split\":\n",
" adata.obs[\"source_split\"] = adata.obs[\"split\"].astype(str)\n",
" adata.obs[\"split\"] = assign_seen_cell_split(adata.obs, RANDOM_SEED)\n",
"elif SPLIT_STRATEGY == \"use_existing\":\n",
" adata.obs[\"split\"] = pd.Categorical(\n",
" adata.obs[\"split\"].astype(str),\n",
" categories=[\"train\", \"val\", \"test\"],\n",
" ordered=True,\n",
" )\n",
"else:\n",
" raise ValueError(f\"Unknown SPLIT_STRATEGY: {SPLIT_STRATEGY}\")\n",
"\n",
"is_control = adata.obs[\"is_control\"].to_numpy(dtype=bool)\n",
"split = adata.obs[\"split\"].astype(str).to_numpy()\n",
"genes = adata.obs[\"perturbation_gene\"].astype(str).to_numpy()\n",
"batches = adata.obs[\"batch_model\"].astype(str).to_numpy()\n",
"\n",
"control_train_idx = np.where(is_control & (split == CONTROL_SPLIT))[0]\n",
"if len(control_train_idx) == 0:\n",
" print(f\"No controls found in CONTROL_SPLIT={CONTROL_SPLIT}; falling back to all controls.\")\n",
" control_train_idx = np.where(is_control)[0]\n",
"if len(control_train_idx) == 0:\n",
" raise ValueError(\"No control cells found.\")\n",
"\n",
"train_conditions = sorted(set(genes[(~is_control) & (split == TRAIN_SPLIT)]))\n",
"val_conditions = sorted(set(genes[(~is_control) & (split == VAL_SPLIT)]))\n",
"eval_conditions = sorted(set(genes[(~is_control) & (split == EVAL_SPLIT)]))\n",
"\n",
"if MAX_TRAIN_CONDITIONS is not None:\n",
" train_conditions = train_conditions[:MAX_TRAIN_CONDITIONS]\n",
"if MAX_EVAL_GENES is not None:\n",
" eval_conditions = eval_conditions[:MAX_EVAL_GENES]\n",
"\n",
"indices_by_split_gene = {split_name: {} for split_name in [TRAIN_SPLIT, VAL_SPLIT, EVAL_SPLIT]}\n",
"for split_name in [TRAIN_SPLIT, VAL_SPLIT, EVAL_SPLIT]:\n",
" split_mask = (~is_control) & (split == split_name)\n",
" allowed = train_conditions if split_name == TRAIN_SPLIT else val_conditions if split_name == VAL_SPLIT else eval_conditions\n",
" for gene in allowed:\n",
" idx = np.where(split_mask & (genes == gene))[0]\n",
" if len(idx) > 0:\n",
" indices_by_split_gene[split_name][gene] = idx\n",
"\n",
"control_idx_by_batch = {}\n",
"for batch_name in sorted(set(batches[control_train_idx])):\n",
" idx = control_train_idx[batches[control_train_idx] == batch_name]\n",
" if len(idx) > 0:\n",
" control_idx_by_batch[batch_name] = idx\n",
"\n",
"gene_vocab = [\"__ctrl__\", \"__unk__\", *train_conditions]\n",
"gene_to_idx = {gene: idx for idx, gene in enumerate(gene_vocab)}\n",
"batch_vocab = [\"__unk__\", *sorted(set(batches))]\n",
"batch_to_idx = {batch: idx for idx, batch in enumerate(batch_vocab)}\n",
"batch_index_per_obs = np.array([batch_to_idx.get(batch, batch_to_idx[\"__unk__\"]) for batch in batches], dtype=np.int64)\n",
"\n",
"control_mean = mean_vector(adata.X[control_train_idx])\n",
"unseen_eval_genes = sorted(set(eval_conditions) - set(train_conditions))\n",
"\n",
"print(\"adata:\", adata)\n",
"print(\"split counts:\")\n",
"print(adata.obs[\"split\"].astype(str).value_counts())\n",
"print(\"train conditions:\", len(train_conditions))\n",
"print(\"val conditions:\", len(val_conditions))\n",
"print(\"eval conditions:\", len(eval_conditions))\n",
"print(\"eval genes missing from training:\", len(unseen_eval_genes))\n",
"if unseen_eval_genes[:10]:\n",
" print(\"first unseen eval genes:\", unseen_eval_genes[:10])\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "63b7c6b8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"LatentDiffusionModel(\n",
" (encoder): MLPEncoder(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=4334, out_features=512, bias=True)\n",
" (1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
" (2): GELU(approximate='none')\n",
" (3): Dropout(p=0.1, inplace=False)\n",
" (4): Linear(in_features=512, out_features=256, bias=True)\n",
" (5): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
" (6): GELU(approximate='none')\n",
" (7): Dropout(p=0.1, inplace=False)\n",
" (8): Linear(in_features=256, out_features=64, bias=True)\n",
" )\n",
" )\n",
" (decoder): MLPDecoder(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=64, out_features=256, bias=True)\n",
" (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
" (2): GELU(approximate='none')\n",
" (3): Dropout(p=0.1, inplace=False)\n",
" (4): Linear(in_features=256, out_features=512, bias=True)\n",
" (5): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
" (6): GELU(approximate='none')\n",
" (7): Dropout(p=0.1, inplace=False)\n",
" (8): Linear(in_features=512, out_features=4334, bias=True)\n",
" )\n",
" )\n",
" (condition_encoder): ConditionEncoder(\n",
" (gene_embedding): Embedding(2058, 64)\n",
" (batch_embedding): Embedding(49, 32)\n",
" (mlp): Sequential(\n",
" (0): Linear(in_features=96, out_features=128, bias=True)\n",
" (1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
" (2): GELU(approximate='none')\n",
" (3): Dropout(p=0.1, inplace=False)\n",
" (4): Linear(in_features=128, out_features=128, bias=True)\n",
" )\n",
" )\n",
" (time_encoder): TimeEmbedding(\n",
" (proj): Sequential(\n",
" (0): Linear(in_features=64, out_features=64, bias=True)\n",
" (1): SiLU()\n",
" (2): Linear(in_features=64, out_features=64, bias=True)\n",
" )\n",
" )\n",
" (denoiser): ConditionalDenoiser(\n",
" (input_proj): Linear(in_features=320, out_features=256, bias=True)\n",
" (blocks): ModuleList(\n",
" (0-3): 4 x ConditionalResidualBlock(\n",
" (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
" (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
" (fc1): Linear(in_features=256, out_features=256, bias=True)\n",
" (fc2): Linear(in_features=256, out_features=256, bias=True)\n",
" (film1): Linear(in_features=192, out_features=512, bias=True)\n",
" (film2): Linear(in_features=192, out_features=512, bias=True)\n",
" (dropout): Dropout(p=0.1, inplace=False)\n",
" )\n",
" )\n",
" (out): Linear(in_features=256, out_features=64, bias=True)\n",
" )\n",
")\n"
]
}
],
"source": [
"class MLPEncoder(nn.Module):\n",
" def __init__(self, input_dim: int, latent_dim: int, hidden_dim: int, dropout: float):\n",
" super().__init__()\n",
" self.net = nn.Sequential(\n",
" nn.Linear(input_dim, hidden_dim * 2),\n",
" nn.LayerNorm(hidden_dim * 2),\n",
" nn.GELU(),\n",
" nn.Dropout(dropout),\n",
" nn.Linear(hidden_dim * 2, hidden_dim),\n",
" nn.LayerNorm(hidden_dim),\n",
" nn.GELU(),\n",
" nn.Dropout(dropout),\n",
" nn.Linear(hidden_dim, latent_dim),\n",
" )\n",
"\n",
" def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
" return self.net(x)\n",
"\n",
"\n",
"class MLPDecoder(nn.Module):\n",
" def __init__(self, latent_dim: int, output_dim: int, hidden_dim: int, dropout: float):\n",
" super().__init__()\n",
" self.net = nn.Sequential(\n",
" nn.Linear(latent_dim, hidden_dim),\n",
" nn.LayerNorm(hidden_dim),\n",
" nn.GELU(),\n",
" nn.Dropout(dropout),\n",
" nn.Linear(hidden_dim, hidden_dim * 2),\n",
" nn.LayerNorm(hidden_dim * 2),\n",
" nn.GELU(),\n",
" nn.Dropout(dropout),\n",
" nn.Linear(hidden_dim * 2, output_dim),\n",
" )\n",
"\n",
" def forward(self, z: torch.Tensor) -> torch.Tensor:\n",
" return F.softplus(self.net(z))\n",
"\n",
"\n",
"class ConditionEncoder(nn.Module):\n",
" def __init__(self, n_genes: int, n_batches: int, condition_dim: int, dropout: float):\n",
" super().__init__()\n",
" gene_dim = max(32, condition_dim // 2)\n",
" batch_dim = max(16, condition_dim // 4)\n",
" self.gene_embedding = nn.Embedding(n_genes, gene_dim)\n",
" self.batch_embedding = nn.Embedding(n_batches, batch_dim)\n",
" self.mlp = nn.Sequential(\n",
" nn.Linear(gene_dim + batch_dim, condition_dim),\n",
" nn.LayerNorm(condition_dim),\n",
" nn.GELU(),\n",
" nn.Dropout(dropout),\n",
" nn.Linear(condition_dim, condition_dim),\n",
" )\n",
"\n",
" def forward(self, gene_idx: torch.Tensor, batch_idx: torch.Tensor) -> torch.Tensor:\n",
" gene_emb = self.gene_embedding(gene_idx)\n",
" batch_emb = self.batch_embedding(batch_idx)\n",
" return self.mlp(torch.cat([gene_emb, batch_emb], dim=-1))\n",
"\n",
"\n",
"class TimeEmbedding(nn.Module):\n",
" def __init__(self, time_dim: int):\n",
" super().__init__()\n",
" self.time_dim = time_dim\n",
" self.proj = nn.Sequential(\n",
" nn.Linear(time_dim, time_dim),\n",
" nn.SiLU(),\n",
" nn.Linear(time_dim, time_dim),\n",
" )\n",
"\n",
" def forward(self, t_index: torch.Tensor) -> torch.Tensor:\n",
" t = t_index.float()\n",
" half = self.time_dim // 2\n",
" freqs = torch.exp(\n",
" -math.log(10000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / max(half, 1)\n",
" )\n",
" angles = t[:, None] * freqs[None, :]\n",
" emb = torch.cat([torch.cos(angles), torch.sin(angles)], dim=-1)\n",
" if emb.shape[1] < self.time_dim:\n",
" emb = F.pad(emb, (0, self.time_dim - emb.shape[1]))\n",
" return self.proj(emb)\n",
"\n",
"\n",
"class ConditionalResidualBlock(nn.Module):\n",
" def __init__(self, hidden_dim: int, cond_dim: int, dropout: float):\n",
" super().__init__()\n",
" self.norm1 = nn.LayerNorm(hidden_dim)\n",
" self.norm2 = nn.LayerNorm(hidden_dim)\n",
" self.fc1 = nn.Linear(hidden_dim, hidden_dim)\n",
" self.fc2 = nn.Linear(hidden_dim, hidden_dim)\n",
" self.film1 = nn.Linear(cond_dim, hidden_dim * 2)\n",
" self.film2 = nn.Linear(cond_dim, hidden_dim * 2)\n",
" self.dropout = nn.Dropout(dropout)\n",
"\n",
" def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:\n",
" gamma1, beta1 = self.film1(cond).chunk(2, dim=-1)\n",
" h = self.norm1(x)\n",
" h = h * (1 + gamma1) + beta1\n",
" h = F.gelu(self.fc1(h))\n",
" h = self.dropout(h)\n",
"\n",
" gamma2, beta2 = self.film2(cond).chunk(2, dim=-1)\n",
" h = self.norm2(h)\n",
" h = h * (1 + gamma2) + beta2\n",
" h = self.fc2(h)\n",
" h = self.dropout(h)\n",
" return x + h\n",
"\n",
"\n",
"class ConditionalDenoiser(nn.Module):\n",
" def __init__(self, latent_dim: int, hidden_dim: int, cond_dim: int, num_blocks: int, dropout: float):\n",
" super().__init__()\n",
" self.input_proj = nn.Linear(latent_dim * 2 + TIME_DIM + CONDITION_DIM, hidden_dim)\n",
" self.blocks = nn.ModuleList(\n",
" [ConditionalResidualBlock(hidden_dim, cond_dim=TIME_DIM + CONDITION_DIM, dropout=dropout) for _ in range(num_blocks)]\n",
" )\n",
" self.out = nn.Linear(hidden_dim, latent_dim)\n",
"\n",
" def forward(self, z_t: torch.Tensor, z_ctrl: torch.Tensor, t_emb: torch.Tensor, cond_emb: torch.Tensor) -> torch.Tensor:\n",
" cond = torch.cat([t_emb, cond_emb], dim=-1)\n",
" h = self.input_proj(torch.cat([z_t, z_ctrl, cond], dim=-1))\n",
" for block in self.blocks:\n",
" h = block(h, cond)\n",
" return self.out(h)\n",
"\n",
"\n",
"class LatentDiffusionModel(nn.Module):\n",
" def __init__(self, input_dim: int, latent_dim: int, hidden_dim: int, condition_dim: int, time_dim: int, n_genes: int, n_batches: int, num_blocks: int, dropout: float):\n",
" super().__init__()\n",
" self.encoder = MLPEncoder(input_dim=input_dim, latent_dim=latent_dim, hidden_dim=hidden_dim, dropout=dropout)\n",
" self.decoder = MLPDecoder(latent_dim=latent_dim, output_dim=input_dim, hidden_dim=hidden_dim, dropout=dropout)\n",
" self.condition_encoder = ConditionEncoder(n_genes=n_genes, n_batches=n_batches, condition_dim=condition_dim, dropout=dropout)\n",
" self.time_encoder = TimeEmbedding(time_dim=time_dim)\n",
" self.denoiser = ConditionalDenoiser(\n",
" latent_dim=latent_dim,\n",
" hidden_dim=hidden_dim,\n",
" cond_dim=time_dim + condition_dim,\n",
" num_blocks=num_blocks,\n",
" dropout=dropout,\n",
" )\n",
"\n",
" def encode_x(self, x: torch.Tensor) -> torch.Tensor:\n",
" return self.encoder(x)\n",
"\n",
" def decode_z(self, z: torch.Tensor) -> torch.Tensor:\n",
" return self.decoder(z)\n",
"\n",
" def reconstruct(self, x: torch.Tensor) -> torch.Tensor:\n",
" return self.decode_z(self.encode_x(x))\n",
"\n",
" def predict_noise(self, z_t: torch.Tensor, t_index: torch.Tensor, z_ctrl: torch.Tensor, gene_idx: torch.Tensor, batch_idx: torch.Tensor) -> torch.Tensor:\n",
" cond_emb = self.condition_encoder(gene_idx, batch_idx)\n",
" t_emb = self.time_encoder(t_index)\n",
" return self.denoiser(z_t=z_t, z_ctrl=z_ctrl, t_emb=t_emb, cond_emb=cond_emb)\n",
"\n",
"\n",
"model = LatentDiffusionModel(\n",
" input_dim=adata.n_vars,\n",
" latent_dim=LATENT_DIM,\n",
" hidden_dim=HIDDEN_DIM,\n",
" condition_dim=CONDITION_DIM,\n",
" time_dim=TIME_DIM,\n",
" n_genes=len(gene_vocab),\n",
" n_batches=len(batch_vocab),\n",
" num_blocks=NUM_DENOISER_BLOCKS,\n",
" dropout=DROPOUT,\n",
").to(DEVICE)\n",
"\n",
"optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)\n",
"\n",
"betas = torch.linspace(BETA_START, BETA_END, DIFFUSION_STEPS, device=DEVICE, dtype=torch.float32)\n",
"alphas = 1.0 - betas\n",
"alpha_bars = torch.cumprod(alphas, dim=0)\n",
"sqrt_alpha_bars = torch.sqrt(alpha_bars)\n",
"sqrt_one_minus_alpha_bars = torch.sqrt(1.0 - alpha_bars)\n",
"sqrt_recip_alphas = torch.sqrt(1.0 / alphas)\n",
"posterior_variance = betas * (1.0 - torch.cat([torch.tensor([1.0], device=DEVICE), alpha_bars[:-1]])) / (1.0 - alpha_bars)\n",
"posterior_variance = posterior_variance.clamp_min(1e-8)\n",
"\n",
"print(model)\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "1705b04b",
"metadata": {},
"outputs": [],
"source": [
"def sample_controls_like_targets(target_idx: np.ndarray, rng: np.random.Generator) -> np.ndarray:\n",
" sampled = []\n",
" for idx in target_idx:\n",
" batch_name = batches[idx]\n",
" pool = control_idx_by_batch.get(batch_name, control_train_idx)\n",
" sampled.append(rng.choice(pool))\n",
" return np.asarray(sampled, dtype=np.int64)\n",
"\n",
"\n",
"def make_condition_batch(condition: str, split_name: str, rng: np.random.Generator) -> dict:\n",
" target_pool = indices_by_split_gene[split_name][condition]\n",
" size = min(BATCH_SIZE, len(target_pool))\n",
" replace = len(target_pool) < size\n",
" target_idx = rng.choice(target_pool, size=size, replace=replace)\n",
" source_idx = sample_controls_like_targets(target_idx, rng)\n",
"\n",
" if MAX_EVAL_CELLS_PER_GENE is not None and len(target_idx) > MAX_EVAL_CELLS_PER_GENE:\n",
" target_idx = target_idx[:MAX_EVAL_CELLS_PER_GENE]\n",
" source_idx = source_idx[:MAX_EVAL_CELLS_PER_GENE]\n",
"\n",
" source_x = torch.from_numpy(take_dense_rows(adata.X, source_idx)).to(DEVICE)\n",
" target_x = torch.from_numpy(take_dense_rows(adata.X, target_idx)).to(DEVICE)\n",
" gene_idx = torch.full(\n",
" (len(target_idx),),\n",
" fill_value=gene_to_idx.get(condition, gene_to_idx[\"__unk__\"]),\n",
" device=DEVICE,\n",
" dtype=torch.long,\n",
" )\n",
" batch_idx = torch.as_tensor(batch_index_per_obs[source_idx], device=DEVICE, dtype=torch.long)\n",
"\n",
" return {\n",
" \"condition\": condition,\n",
" \"source_idx\": source_idx,\n",
" \"target_idx\": target_idx,\n",
" \"source_x\": source_x,\n",
" \"target_x\": target_x,\n",
" \"gene_idx\": gene_idx,\n",
" \"batch_idx\": batch_idx,\n",
" }\n",
"\n",
"\n",
"def q_sample(z_0: torch.Tensor, t_index: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:\n",
" sqrt_ab = sqrt_alpha_bars.index_select(0, t_index).unsqueeze(-1)\n",
" sqrt_omb = sqrt_one_minus_alpha_bars.index_select(0, t_index).unsqueeze(-1)\n",
" return sqrt_ab * z_0 + sqrt_omb * noise\n",
"\n",
"\n",
"def compute_loss_terms(batch: dict) -> dict[str, torch.Tensor]:\n",
" z_ctrl = model.encode_x(batch[\"source_x\"])\n",
" z_target = model.encode_x(batch[\"target_x\"])\n",
" z_target_bar = sinkhorn_barycentric_targets(\n",
" x0=z_ctrl,\n",
" x1=z_target,\n",
" epsilon=SINKHORN_EPSILON,\n",
" n_iters=SINKHORN_ITERS,\n",
" )\n",
"\n",
" batch_size = z_target_bar.shape[0]\n",
" t_index = torch.randint(0, DIFFUSION_STEPS, (batch_size,), device=DEVICE, dtype=torch.long)\n",
" noise = torch.randn_like(z_target_bar)\n",
" z_t = q_sample(z_target_bar, t_index, noise)\n",
"\n",
" pred_noise = model.predict_noise(\n",
" z_t=z_t,\n",
" t_index=t_index,\n",
" z_ctrl=z_ctrl,\n",
" gene_idx=batch[\"gene_idx\"],\n",
" batch_idx=batch[\"batch_idx\"],\n",
" )\n",
"\n",
" diffusion_loss = F.mse_loss(pred_noise, noise)\n",
"\n",
" z0_hat = (z_t - sqrt_one_minus_alpha_bars.index_select(0, t_index).unsqueeze(-1) * pred_noise) / sqrt_alpha_bars.index_select(0, t_index).unsqueeze(-1)\n",
" pred_x1 = model.decode_z(z0_hat)\n",
"\n",
" source_recon = model.reconstruct(batch[\"source_x\"])\n",
" target_recon = model.reconstruct(batch[\"target_x\"])\n",
"\n",
" recon_loss = F.mse_loss(source_recon, batch[\"source_x\"]) + F.mse_loss(target_recon, batch[\"target_x\"])\n",
" endpoint_loss = F.mse_loss(z0_hat, z_target_bar)\n",
" mmd_loss = rbf_mmd_loss(pred_x1, batch[\"target_x\"])\n",
" mean_loss = F.mse_loss(pred_x1.mean(dim=0), batch[\"target_x\"].mean(dim=0))\n",
"\n",
" total_loss = (\n",
" RECON_WEIGHT * recon_loss\n",
" + DIFFUSION_WEIGHT * diffusion_loss\n",
" + ENDPOINT_WEIGHT * endpoint_loss\n",
" + MMD_WEIGHT * mmd_loss\n",
" + MEAN_WEIGHT * mean_loss\n",
" )\n",
"\n",
" return {\n",
" \"total\": total_loss,\n",
" \"recon\": recon_loss.detach(),\n",
" \"diffusion\": diffusion_loss.detach(),\n",
" \"endpoint\": endpoint_loss.detach(),\n",
" \"mmd\": mmd_loss.detach(),\n",
" \"mean\": mean_loss.detach(),\n",
" }\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def evaluate_split_loss(split_name: str, conditions: list[str], seed: int, steps: int) -> dict[str, float]:\n",
" if not conditions:\n",
" return {}\n",
"\n",
" model.eval()\n",
" rng = np.random.default_rng(seed)\n",
" subset = conditions.copy()\n",
" rng.shuffle(subset)\n",
" subset = subset[: min(len(subset), steps)]\n",
"\n",
" rows = []\n",
" for condition in subset:\n",
" batch = make_condition_batch(condition, split_name=split_name, rng=rng)\n",
" losses = compute_loss_terms(batch)\n",
" rows.append({key: float(value.item()) for key, value in losses.items()})\n",
"\n",
" return pd.DataFrame(rows).mean().to_dict()\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def sample_latent(z_ctrl: torch.Tensor, gene_idx: torch.Tensor, batch_idx: torch.Tensor) -> torch.Tensor:\n",
" z = torch.randn_like(z_ctrl)\n",
" for t in reversed(range(DIFFUSION_STEPS)):\n",
" t_index = torch.full((z.shape[0],), t, device=DEVICE, dtype=torch.long)\n",
" pred_noise = model.predict_noise(\n",
" z_t=z,\n",
" t_index=t_index,\n",
" z_ctrl=z_ctrl,\n",
" gene_idx=gene_idx,\n",
" batch_idx=batch_idx,\n",
" )\n",
"\n",
" beta_t = betas[t]\n",
" sqrt_one_minus_ab_t = sqrt_one_minus_alpha_bars[t]\n",
" sqrt_recip_alpha_t = sqrt_recip_alphas[t]\n",
"\n",
" model_mean = sqrt_recip_alpha_t * (z - beta_t / sqrt_one_minus_ab_t * pred_noise)\n",
" if t > 0:\n",
" noise = torch.randn_like(z)\n",
" z = model_mean + torch.sqrt(posterior_variance[t]) * noise\n",
" else:\n",
" z = model_mean\n",
" return z\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def predict_gene_expression(condition: str, split_name: str, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray, pd.DataFrame]:\n",
" target_idx_all = indices_by_split_gene[split_name][condition]\n",
" if MAX_EVAL_CELLS_PER_GENE is not None and len(target_idx_all) > MAX_EVAL_CELLS_PER_GENE:\n",
" target_idx_all = target_idx_all[:MAX_EVAL_CELLS_PER_GENE]\n",
"\n",
" pred_blocks = []\n",
" real_blocks = []\n",
" obs_blocks = []\n",
" gene_id = gene_to_idx.get(condition, gene_to_idx[\"__unk__\"])\n",
" condition_source = \"trained_gene_embedding\" if condition in gene_to_idx else \"unk_gene_embedding\"\n",
"\n",
" for start in range(0, len(target_idx_all), PRED_BATCH_SIZE):\n",
" target_idx = target_idx_all[start : start + PRED_BATCH_SIZE]\n",
" source_idx = sample_controls_like_targets(target_idx, rng)\n",
"\n",
" source_x = torch.from_numpy(take_dense_rows(adata.X, source_idx)).to(DEVICE)\n",
" z_ctrl = model.encode_x(source_x)\n",
" batch_idx = torch.as_tensor(batch_index_per_obs[source_idx], device=DEVICE, dtype=torch.long)\n",
" gene_idx = torch.full((len(target_idx),), gene_id, device=DEVICE, dtype=torch.long)\n",
"\n",
" z_hat = sample_latent(z_ctrl=z_ctrl, gene_idx=gene_idx, batch_idx=batch_idx)\n",
" pred_x = model.decode_z(z_hat).cpu().numpy().astype(np.float32)\n",
" real_x = take_dense_rows(adata.X, target_idx)\n",
"\n",
" pred_blocks.append(pred_x)\n",
" real_blocks.append(real_x)\n",
"\n",
" obs_chunk = adata.obs.iloc[target_idx][[\"perturbation_gene\", \"condition\", \"split\"]].copy()\n",
" obs_chunk[\"target_cell_id\"] = adata.obs_names[target_idx].astype(str)\n",
" obs_chunk[\"sampled_control_cell_id\"] = adata.obs_names[source_idx].astype(str)\n",
" obs_chunk[\"condition_source\"] = condition_source\n",
" obs_blocks.append(obs_chunk)\n",
"\n",
" pred_gene = np.vstack(pred_blocks).astype(np.float32)\n",
" real_gene = np.vstack(real_blocks).astype(np.float32)\n",
" obs_gene = pd.concat(obs_blocks, axis=0)\n",
" return pred_gene, real_gene, obs_gene\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def evaluate_gene_metrics(split_name: str, conditions: list[str], seed: int) -> tuple[pd.DataFrame, ad.AnnData, ad.AnnData]:\n",
" rng = np.random.default_rng(seed)\n",
"\n",
" pred_blocks = []\n",
" real_blocks = []\n",
" obs_blocks = []\n",
" metrics_rows = []\n",
"\n",
" for condition in conditions:\n",
" pred_gene, real_gene, obs_gene = predict_gene_expression(condition=condition, split_name=split_name, rng=rng)\n",
" pred_blocks.append(pred_gene)\n",
" real_blocks.append(real_gene)\n",
" obs_blocks.append(obs_gene)\n",
"\n",
" pred_mean = pred_gene.mean(axis=0)\n",
" real_mean = real_gene.mean(axis=0)\n",
" pred_delta = pred_mean - control_mean\n",
" real_delta = real_mean - control_mean\n",
"\n",
" metrics_rows.append(\n",
" {\n",
" \"perturbation_gene\": condition,\n",
" \"n_cells\": int(real_gene.shape[0]),\n",
" \"mse_mean\": float(np.mean((pred_mean - real_mean) ** 2)),\n",
" \"mae_mean\": float(np.mean(np.abs(pred_mean - real_mean))),\n",
" \"pearson_mean\": safe_pearson(pred_mean, real_mean),\n",
" \"pearson_delta\": safe_pearson(pred_delta, real_delta),\n",
" \"delta_l2\": float(np.linalg.norm(pred_delta - real_delta)),\n",
" }\n",
" )\n",
"\n",
" pred = ad.AnnData(\n",
" X=np.vstack(pred_blocks).astype(np.float32),\n",
" obs=pd.concat(obs_blocks, axis=0),\n",
" var=adata.var.copy(),\n",
" )\n",
" real = ad.AnnData(\n",
" X=np.vstack(real_blocks).astype(np.float32),\n",
" obs=pd.concat(obs_blocks, axis=0),\n",
" var=adata.var.copy(),\n",
" )\n",
" metrics = pd.DataFrame(metrics_rows).sort_values(\"perturbation_gene\").reset_index(drop=True)\n",
" return metrics, pred, real\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "191e2bd7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'epoch': 1, 'train_total': 1.0926, 'train_diffusion': 0.7012, 'train_recon': 0.2978, 'val_total': 0.565}\n",
"{'epoch': 2, 'train_total': 0.5127, 'train_diffusion': 0.2037, 'train_recon': 0.2779, 'val_total': 0.4346}\n",
"{'epoch': 3, 'train_total': 0.4294, 'train_diffusion': 0.1309, 'train_recon': 0.274, 'val_total': 0.4085}\n",
"{'epoch': 4, 'train_total': 0.394, 'train_diffusion': 0.0996, 'train_recon': 0.2738, 'val_total': 0.3641}\n",
"{'epoch': 5, 'train_total': 0.3757, 'train_diffusion': 0.0805, 'train_recon': 0.2757, 'val_total': 0.36}\n",
"{'epoch': 6, 'train_total': 0.3575, 'train_diffusion': 0.0619, 'train_recon': 0.2768, 'val_total': 0.3343}\n",
"{'epoch': 7, 'train_total': 0.3451, 'train_diffusion': 0.0544, 'train_recon': 0.2732, 'val_total': 0.3495}\n",
"{'epoch': 8, 'train_total': 0.3357, 'train_diffusion': 0.0471, 'train_recon': 0.2722, 'val_total': 0.3428}\n",
"{'epoch': 9, 'train_total': 0.331, 'train_diffusion': 0.0423, 'train_recon': 0.2719, 'val_total': 0.335}\n",
"{'epoch': 10, 'train_total': 0.3268, 'train_diffusion': 0.0372, 'train_recon': 0.273, 'val_total': 0.358}\n",
"{'epoch': 11, 'train_total': 0.3188, 'train_diffusion': 0.0337, 'train_recon': 0.2695, 'val_total': 0.3339}\n",
"{'epoch': 12, 'train_total': 0.3157, 'train_diffusion': 0.0301, 'train_recon': 0.2707, 'val_total': 0.3189}\n",
"{'epoch': 13, 'train_total': 0.3117, 'train_diffusion': 0.0274, 'train_recon': 0.2701, 'val_total': 0.3391}\n",
"{'epoch': 14, 'train_total': 0.3093, 'train_diffusion': 0.0268, 'train_recon': 0.2684, 'val_total': 0.339}\n",
"{'epoch': 15, 'train_total': 0.3056, 'train_diffusion': 0.0263, 'train_recon': 0.2664, 'val_total': 0.3072}\n",
"{'epoch': 16, 'train_total': 0.3104, 'train_diffusion': 0.0275, 'train_recon': 0.2692, 'val_total': 0.3288}\n",
"{'epoch': 17, 'train_total': 0.3063, 'train_diffusion': 0.0268, 'train_recon': 0.2665, 'val_total': 0.3394}\n",
"{'epoch': 18, 'train_total': 0.3053, 'train_diffusion': 0.0243, 'train_recon': 0.2675, 'val_total': 0.3203}\n",
"{'epoch': 19, 'train_total': 0.3035, 'train_diffusion': 0.0248, 'train_recon': 0.266, 'val_total': 0.3337}\n",
"{'epoch': 20, 'train_total': 0.304, 'train_diffusion': 0.0268, 'train_recon': 0.2633, 'val_total': 0.3275}\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" epoch | \n",
" train_total | \n",
" train_recon | \n",
" train_diffusion | \n",
" train_endpoint | \n",
" train_mmd | \n",
" train_mean | \n",
" val_total | \n",
" val_recon | \n",
" val_diffusion | \n",
" val_endpoint | \n",
" val_mmd | \n",
" val_mean | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 1 | \n",
" 1.092597 | \n",
" 0.297794 | \n",
" 0.701235 | \n",
" 0.123018 | \n",
" 0.143856 | \n",
" 0.017673 | \n",
" 0.564962 | \n",
" 0.276370 | \n",
" 0.222312 | \n",
" 0.028087 | \n",
" 0.201846 | \n",
" 0.032051 | \n",
"
\n",
" \n",
" | 1 | \n",
" 2 | \n",
" 0.512674 | \n",
" 0.277870 | \n",
" 0.203705 | \n",
" 0.025684 | \n",
" 0.106283 | \n",
" 0.007628 | \n",
" 0.434584 | \n",
" 0.273373 | \n",
" 0.105633 | \n",
" 0.010058 | \n",
" 0.171663 | \n",
" 0.033383 | \n",
"
\n",
" \n",
" | 2 | \n",
" 3 | \n",
" 0.429386 | \n",
" 0.273969 | \n",
" 0.130892 | \n",
" 0.016188 | \n",
" 0.088635 | \n",
" 0.007567 | \n",
" 0.408545 | \n",
" 0.279798 | \n",
" 0.084368 | \n",
" 0.007179 | \n",
" 0.154740 | \n",
" 0.025315 | \n",
"
\n",
" \n",
" | 3 | \n",
" 4 | \n",
" 0.394016 | \n",
" 0.273814 | \n",
" 0.099629 | \n",
" 0.011507 | \n",
" 0.077374 | \n",
" 0.007082 | \n",
" 0.364078 | \n",
" 0.274999 | \n",
" 0.035762 | \n",
" 0.003864 | \n",
" 0.178134 | \n",
" 0.033572 | \n",
"
\n",
" \n",
" | 4 | \n",
" 5 | \n",
" 0.375732 | \n",
" 0.275702 | \n",
" 0.080474 | \n",
" 0.007924 | \n",
" 0.077091 | \n",
" 0.007886 | \n",
" 0.359967 | \n",
" 0.271042 | \n",
" 0.042826 | \n",
" 0.004512 | \n",
" 0.167649 | \n",
" 0.027078 | \n",
"
\n",
" \n",
" | 5 | \n",
" 6 | \n",
" 0.357506 | \n",
" 0.276800 | \n",
" 0.061926 | \n",
" 0.005487 | \n",
" 0.084314 | \n",
" 0.007605 | \n",
" 0.334270 | \n",
" 0.270256 | \n",
" 0.025244 | \n",
" 0.003013 | \n",
" 0.115518 | \n",
" 0.025712 | \n",
"
\n",
" \n",
" | 6 | \n",
" 7 | \n",
" 0.345125 | \n",
" 0.273160 | \n",
" 0.054362 | \n",
" 0.004331 | \n",
" 0.079824 | \n",
" 0.007456 | \n",
" 0.349526 | \n",
" 0.271062 | \n",
" 0.035995 | \n",
" 0.002528 | \n",
" 0.137571 | \n",
" 0.027447 | \n",
"
\n",
" \n",
" | 7 | \n",
" 8 | \n",
" 0.335657 | \n",
" 0.272180 | \n",
" 0.047059 | \n",
" 0.003757 | \n",
" 0.069465 | \n",
" 0.007593 | \n",
" 0.342845 | \n",
" 0.268517 | \n",
" 0.019302 | \n",
" 0.001905 | \n",
" 0.151532 | \n",
" 0.038920 | \n",
"
\n",
" \n",
" | 8 | \n",
" 9 | \n",
" 0.330977 | \n",
" 0.271936 | \n",
" 0.042338 | \n",
" 0.003208 | \n",
" 0.072519 | \n",
" 0.007846 | \n",
" 0.334984 | \n",
" 0.264954 | \n",
" 0.018661 | \n",
" 0.001880 | \n",
" 0.138549 | \n",
" 0.036574 | \n",
"
\n",
" \n",
" | 9 | \n",
" 10 | \n",
" 0.326805 | \n",
" 0.273013 | \n",
" 0.037230 | \n",
" 0.003037 | \n",
" 0.069798 | \n",
" 0.008064 | \n",
" 0.357998 | \n",
" 0.273934 | \n",
" 0.021520 | \n",
" 0.003141 | \n",
" 0.186460 | \n",
" 0.042328 | \n",
"
\n",
" \n",
" | 10 | \n",
" 11 | \n",
" 0.318764 | \n",
" 0.269548 | \n",
" 0.033744 | \n",
" 0.002867 | \n",
" 0.064630 | \n",
" 0.007575 | \n",
" 0.333905 | \n",
" 0.264447 | \n",
" 0.028742 | \n",
" 0.002804 | \n",
" 0.133314 | \n",
" 0.025983 | \n",
"
\n",
" \n",
" | 11 | \n",
" 12 | \n",
" 0.315714 | \n",
" 0.270734 | \n",
" 0.030088 | \n",
" 0.002676 | \n",
" 0.058792 | \n",
" 0.007674 | \n",
" 0.318947 | \n",
" 0.261437 | \n",
" 0.016157 | \n",
" 0.001626 | \n",
" 0.138025 | \n",
" 0.026738 | \n",
"
\n",
" \n",
" | 12 | \n",
" 13 | \n",
" 0.311701 | \n",
" 0.270063 | \n",
" 0.027354 | \n",
" 0.002231 | \n",
" 0.059132 | \n",
" 0.007254 | \n",
" 0.339054 | \n",
" 0.263594 | \n",
" 0.012227 | \n",
" 0.001638 | \n",
" 0.141034 | \n",
" 0.048310 | \n",
"
\n",
" \n",
" | 13 | \n",
" 14 | \n",
" 0.309325 | \n",
" 0.268446 | \n",
" 0.026806 | \n",
" 0.002427 | \n",
" 0.059312 | \n",
" 0.006929 | \n",
" 0.339031 | \n",
" 0.273432 | \n",
" 0.013024 | \n",
" 0.001383 | \n",
" 0.117266 | \n",
" 0.040157 | \n",
"
\n",
" \n",
" | 14 | \n",
" 15 | \n",
" 0.305622 | \n",
" 0.266382 | \n",
" 0.026251 | \n",
" 0.002382 | \n",
" 0.052757 | \n",
" 0.006523 | \n",
" 0.307164 | \n",
" 0.259813 | \n",
" 0.013409 | \n",
" 0.001333 | \n",
" 0.112954 | \n",
" 0.021979 | \n",
"
\n",
" \n",
" | 15 | \n",
" 16 | \n",
" 0.310393 | \n",
" 0.269180 | \n",
" 0.027466 | \n",
" 0.002505 | \n",
" 0.052884 | \n",
" 0.007206 | \n",
" 0.328840 | \n",
" 0.269686 | \n",
" 0.021865 | \n",
" 0.003770 | \n",
" 0.106170 | \n",
" 0.024787 | \n",
"
\n",
" \n",
" | 16 | \n",
" 17 | \n",
" 0.306289 | \n",
" 0.266506 | \n",
" 0.026796 | \n",
" 0.002581 | \n",
" 0.051856 | \n",
" 0.006512 | \n",
" 0.339396 | \n",
" 0.269045 | \n",
" 0.014996 | \n",
" 0.002078 | \n",
" 0.153371 | \n",
" 0.038978 | \n",
"
\n",
" \n",
" | 17 | \n",
" 18 | \n",
" 0.305292 | \n",
" 0.267518 | \n",
" 0.024289 | \n",
" 0.002265 | \n",
" 0.054927 | \n",
" 0.006859 | \n",
" 0.320261 | \n",
" 0.265096 | \n",
" 0.015183 | \n",
" 0.002067 | \n",
" 0.129587 | \n",
" 0.025990 | \n",
"
\n",
" \n",
" | 18 | \n",
" 19 | \n",
" 0.303483 | \n",
" 0.265995 | \n",
" 0.024794 | \n",
" 0.002212 | \n",
" 0.051210 | \n",
" 0.006467 | \n",
" 0.333737 | \n",
" 0.272283 | \n",
" 0.016100 | \n",
" 0.002361 | \n",
" 0.113277 | \n",
" 0.032845 | \n",
"
\n",
" \n",
" | 19 | \n",
" 20 | \n",
" 0.303973 | \n",
" 0.263291 | \n",
" 0.026823 | \n",
" 0.002629 | \n",
" 0.055363 | \n",
" 0.007008 | \n",
" 0.327471 | \n",
" 0.274065 | \n",
" 0.014217 | \n",
" 0.002168 | \n",
" 0.141412 | \n",
" 0.023963 | \n",
"
\n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"_fig_tag = f\"{SPLIT_STRATEGY}_{EVAL_SPLIT}\"\n",
"plot_training_history(history, tag=_fig_tag)\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "9e0ba50e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Aggregate summary:\n",
"method_name conditional_diffusion_method\n",
"input_path /root/final/data/processed/scPerturb/predictio...\n",
"eval_split test\n",
"train_split train\n",
"val_split val\n",
"n_control_cells_train 8653\n",
"n_eval_cells 29965\n",
"n_train_genes 2056\n",
"n_eval_genes 2056\n",
"n_eval_genes_missing_from_train 0\n",
"mse_mean_avg 0.02029\n",
"mae_mean_avg 0.09986\n",
"pearson_mean_avg 0.974649\n",
"pearson_delta_avg 0.061978\n",
"delta_l2_avg 8.738639\n",
"dtype: object\n",
"\n",
"Saved:\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/conditional_diffusion_model.pt\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/pred_conditional_diffusion_method_seen_cell_split_test.h5ad\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/real_conditional_diffusion_method_seen_cell_split_test.h5ad\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/metrics_by_gene_conditional_diffusion_method_seen_cell_split_test.csv\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/aggregate_conditional_diffusion_method_seen_cell_split_test.csv\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/summary_conditional_diffusion_method_seen_cell_split_test.csv\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/training_history_conditional_diffusion_method_seen_cell_split_test.csv\n",
"- /root/final/baseline/outputs/conditional_diffusion_method/config_conditional_diffusion_method_seen_cell_split_test.json\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" perturbation_gene | \n",
" n_cells | \n",
" mse_mean | \n",
" mae_mean | \n",
" pearson_mean | \n",
" pearson_delta | \n",
" delta_l2 | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" AAAS | \n",
" 21 | \n",
" 0.006507 | \n",
" 0.061623 | \n",
" 0.991500 | \n",
" 0.036507 | \n",
" 5.310577 | \n",
"
\n",
" \n",
" | 1 | \n",
" AAMP | \n",
" 10 | \n",
" 0.024897 | \n",
" 0.116168 | \n",
" 0.970339 | \n",
" -0.179578 | \n",
" 10.387674 | \n",
"
\n",
" \n",
" | 2 | \n",
" AARS | \n",
" 4 | \n",
" 0.044199 | \n",
" 0.159510 | \n",
" 0.944223 | \n",
" 0.299598 | \n",
" 13.840421 | \n",
"
\n",
" \n",
" | 3 | \n",
" AARS2 | \n",
" 18 | \n",
" 0.008631 | \n",
" 0.071208 | \n",
" 0.988610 | \n",
" -0.082087 | \n",
" 6.115983 | \n",
"
\n",
" \n",
" | 4 | \n",
" AASDHPPT | \n",
" 5 | \n",
" 0.032459 | \n",
" 0.137356 | \n",
" 0.960420 | \n",
" 0.075441 | \n",
" 11.860708 | \n",
"
\n",
" \n",
"
\n",
"