From 4cf258587a251c0b71bf9db6d617a86fcb8da773 Mon Sep 17 00:00:00 2001 From: Silja Ohlsen Date: Sun, 5 Jul 2026 07:22:59 +0000 Subject: [PATCH] Fix device mismatch in MIOFlow core and upgrade EBD tutorial to GAGA autoencoder pipeline This commit fixes a critical GPU device mismatch in the MIOFlow core and massively overhauls the EBD tutorial notebook to use the modern Geometric Autoencoder (GAGA) pipeline instead of the older, manual PHATE DataFrame approach. ### Core Changes (mioflow/mioflow.py): * Fixed GPU Device Mismatch Bug: When running on CUDA, _encode() and decode_to_gene_space() previously crashed with a device mismatch (cuda:0 vs cpu). The internal torch.tensor(X_scaled) and torch.tensor(traj_flat) calls were defaulting to the CPU, while the gaga_autoencoder model sat on the GPU. Added device=self.device to these tensor instantiations to ensure device synchronization. ### Tutorial Notebook Changes (tutorials/02_EBD_preprocessing_and_vanilla_colab.ipynb): * Removed legacy PHATE DataFrame logic and instead set discrete_time directly via pd.factorize. * Integrated GAGA Autoencoder Training: Added the missing cell to explicitly train the GAGA autoencoder (fit_gaga) using X_pca and X_phate. * Corrected MIOFlow Initialization: Replaced the deprecated input_df and obsm_key parameters with the modernized Autoencoder configuration (gaga_model=gaga_model and gaga_input_key='X_pca'). * Overhauled Loss Visualization: Replaced the generic plot_losses() function call with a detailed, custom matplotlib block. * Modernized Trajectory Visualization (Step 11): Re-wrote the visualization block to natively evaluate background cells through the trained gaga_model.encode() pipeline, plotting the trajectories directly on the GAGA latent embedding instead of PHATE. * Fixed Scanpy Namespace Collision: Added import scanpy as sc directly before running highly_variable_genes. The previous scatterplot assignment (sc = ax.scatter) overwrote the scanpy as sc module. * Updated Single-Trajectory Selection: Swapped the background scatter data from true_data (PHATE) to gaga_embedding (GAGA). --- mioflow/mioflow.py | 4 +- ..._EBD_preprocessing_and_vanilla_colab.ipynb | 435 ++++++++++++++---- 2 files changed, 352 insertions(+), 87 deletions(-) diff --git a/mioflow/mioflow.py b/mioflow/mioflow.py index ff8893c..557fd81 100644 --- a/mioflow/mioflow.py +++ b/mioflow/mioflow.py @@ -230,7 +230,7 @@ def _encode(self) -> Tuple[np.ndarray, np.ndarray]: self.gaga_autoencoder.eval() with torch.no_grad(): embedding = self.gaga_autoencoder.encode( - torch.tensor(X_scaled) + torch.tensor(X_scaled, device=self.device) ).cpu().numpy().astype(np.float64) else: embedding = X_raw.astype(np.float64) @@ -362,7 +362,7 @@ def decode_to_gene_space(self) -> np.ndarray: self.gaga_autoencoder.eval() with torch.no_grad(): traj_pca_gaga = self.gaga_autoencoder.decode( - torch.tensor(traj_flat, dtype=torch.float32) + torch.tensor(traj_flat, dtype=torch.float32, device=self.device) ).cpu().numpy() # Inverse-scale back to PCA space diff --git a/tutorials/02_EBD_preprocessing_and_vanilla_colab.ipynb b/tutorials/02_EBD_preprocessing_and_vanilla_colab.ipynb index 02e7dfa..9346628 100644 --- a/tutorials/02_EBD_preprocessing_and_vanilla_colab.ipynb +++ b/tutorials/02_EBD_preprocessing_and_vanilla_colab.ipynb @@ -1,20 +1,4 @@ { - "nbformat": 4, - "nbformat_minor": 5, - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.10.0" - }, - "colab": { - "provenance": [] - } - }, "cells": [ { "cell_type": "markdown", @@ -40,7 +24,11 @@ "cell_type": "markdown", "id": "install-md", "metadata": {}, - "source": "## 0. Installation\n\nRun the cell below once to install all required packages." + "source": [ + "## 0. Installation\n", + "\n", + "Run the cell below once to install all required packages." + ] }, { "cell_type": "code", @@ -48,7 +36,9 @@ "id": "install", "metadata": {}, "outputs": [], - "source": "!pip install mioflow phate" + "source": [ + "!pip install mioflow phate" + ] }, { "cell_type": "code", @@ -82,7 +72,29 @@ "id": "imports", "metadata": {}, "outputs": [], - "source": "import os\nimport urllib.request\nimport zipfile\nimport shutil\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nimport scanpy as sc\nimport phate\nimport torch\n\nfrom MIOFlow.mioflow import MIOFlow\nfrom MIOFlow.plots import plot_losses\nfrom MIOFlow.utils import set_seeds\n\nset_seeds(0)\n\nuse_cuda = torch.cuda.is_available()\nprint(f'Using CUDA: {use_cuda}')" + "source": [ + "import os\n", + "import urllib.request\n", + "import zipfile\n", + "import shutil\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import matplotlib.pyplot as plt\n", + "import seaborn as sns\n", + "\n", + "import scanpy as sc\n", + "import phate\n", + "import torch\n", + "\n", + "from mioflow.mioflow import MIOFlow\n", + "\n", + "torch.manual_seed(0)\n", + "np.random.seed(0)\n", + "\n", + "use_cuda = torch.cuda.is_available()\n", + "print(f'Using CUDA: {use_cuda}')" + ] }, { "cell_type": "markdown", @@ -112,19 +124,88 @@ "id": "download", "metadata": {}, "outputs": [], - "source": "RAW_DATA_DIR = 'data/raw/scRNAseq'\nEXPECTED_DIRS = ['T0_1A', 'T2_3B', 'T4_5C', 'T6_7D', 'T8_9E']\n\ndef data_already_downloaded(base_dir, expected_dirs):\n return all(os.path.isdir(os.path.join(base_dir, d)) for d in expected_dirs)\n\nif data_already_downloaded(RAW_DATA_DIR, EXPECTED_DIRS):\n print('Data already present — skipping download.')\nelse:\n os.makedirs(RAW_DATA_DIR, exist_ok=True)\n\n url = 'https://data.mendeley.com/public-api/zip/v6n743h5ng/download/1'\n zip_file = os.path.join(RAW_DATA_DIR, 'v6n743h5ng-1.zip')\n\n print(f'Downloading from {url} ...')\n req = urllib.request.Request(url, headers={\n 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'\n })\n with urllib.request.urlopen(req) as response, open(zip_file, 'wb') as out_file:\n shutil.copyfileobj(response, out_file)\n\n # Extract outer zip to a temp directory\n temp_extract = os.path.join(RAW_DATA_DIR, 'temp_extract')\n os.makedirs(temp_extract, exist_ok=True)\n\n print('Extracting ...')\n with zipfile.ZipFile(zip_file, 'r') as zip_ref:\n zip_ref.extractall(temp_extract)\n\n # Extract inner scRNAseq.zip if present\n scrna_zip = os.path.join(temp_extract, 'scRNAseq.zip')\n if os.path.exists(scrna_zip):\n with zipfile.ZipFile(scrna_zip, 'r') as zip_ref:\n zip_ref.extractall(temp_extract)\n\n # Move sample directories to RAW_DATA_DIR\n scrna_folder = os.path.join(temp_extract, 'scRNAseq')\n src_root = scrna_folder if os.path.exists(scrna_folder) else temp_extract\n for item in EXPECTED_DIRS:\n src = os.path.join(src_root, item)\n dst = os.path.join(RAW_DATA_DIR, item)\n if os.path.exists(src):\n shutil.move(src, dst)\n\n shutil.rmtree(temp_extract)\n os.remove(zip_file)\n print(f'Done. Data extracted to {RAW_DATA_DIR}')\n\nprint('Data directories found:')\nfor d in EXPECTED_DIRS:\n path = os.path.join(RAW_DATA_DIR, d)\n status = 'OK' if os.path.isdir(path) else 'MISSING'\n print(f' {path} [{status}]')" + "source": [ + "RAW_DATA_DIR = 'data/raw/scRNAseq'\n", + "EXPECTED_DIRS = ['T0_1A', 'T2_3B', 'T4_5C', 'T6_7D', 'T8_9E']\n", + "\n", + "def data_already_downloaded(base_dir, expected_dirs):\n", + " return all(os.path.isdir(os.path.join(base_dir, d)) for d in expected_dirs)\n", + "\n", + "if data_already_downloaded(RAW_DATA_DIR, EXPECTED_DIRS):\n", + " print('Data already present — skipping download.')\n", + "else:\n", + " os.makedirs(RAW_DATA_DIR, exist_ok=True)\n", + "\n", + " url = 'https://data.mendeley.com/public-api/zip/v6n743h5ng/download/1'\n", + " zip_file = os.path.join(RAW_DATA_DIR, 'v6n743h5ng-1.zip')\n", + "\n", + " print(f'Downloading from {url} ...')\n", + " req = urllib.request.Request(url, headers={\n", + " 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'\n", + " })\n", + " with urllib.request.urlopen(req) as response, open(zip_file, 'wb') as out_file:\n", + " shutil.copyfileobj(response, out_file)\n", + "\n", + " # Extract outer zip to a temp directory\n", + " temp_extract = os.path.join(RAW_DATA_DIR, 'temp_extract')\n", + " os.makedirs(temp_extract, exist_ok=True)\n", + "\n", + " print('Extracting ...')\n", + " with zipfile.ZipFile(zip_file, 'r') as zip_ref:\n", + " zip_ref.extractall(temp_extract)\n", + "\n", + " # Extract inner scRNAseq.zip if present\n", + " scrna_zip = os.path.join(temp_extract, 'scRNAseq.zip')\n", + " if os.path.exists(scrna_zip):\n", + " with zipfile.ZipFile(scrna_zip, 'r') as zip_ref:\n", + " zip_ref.extractall(temp_extract)\n", + "\n", + " # Move sample directories to RAW_DATA_DIR\n", + " scrna_folder = os.path.join(temp_extract, 'scRNAseq')\n", + " src_root = scrna_folder if os.path.exists(scrna_folder) else temp_extract\n", + " for item in EXPECTED_DIRS:\n", + " src = os.path.join(src_root, item)\n", + " dst = os.path.join(RAW_DATA_DIR, item)\n", + " if os.path.exists(src):\n", + " shutil.move(src, dst)\n", + "\n", + " shutil.rmtree(temp_extract)\n", + " os.remove(zip_file)\n", + " print(f'Done. Data extracted to {RAW_DATA_DIR}')\n", + "\n", + "print('Data directories found:')\n", + "for d in EXPECTED_DIRS:\n", + " path = os.path.join(RAW_DATA_DIR, d)\n", + " status = 'OK' if os.path.isdir(path) else 'MISSING'\n", + " print(f' {path} [{status}]')" + ] }, { "cell_type": "markdown", "id": "preprocessing-md", "metadata": {}, - "source": "## 3. Preprocessing\n\nWe follow the standard scRNA-seq preprocessing pipeline:\n\n1. Load 10X data with `scanpy` and concatenate all time points (with QC metrics)\n2. Filter cells by library size (remove top and bottom 20% per sample)\n3. Remove genes expressed in fewer than 10 cells\n4. Library-size normalise\n5. Remove dead cells (high mitochondrial RNA expression)\n6. Square-root transform" + "source": [ + "## 3. Preprocessing\n", + "\n", + "We follow the standard scRNA-seq preprocessing pipeline:\n", + "\n", + "1. Load 10X data with `scanpy` and concatenate all time points (with QC metrics)\n", + "2. Filter cells by library size (remove top and bottom 20% per sample)\n", + "3. Remove genes expressed in fewer than 10 cells\n", + "4. Library-size normalise\n", + "5. Remove dead cells (high mitochondrial RNA expression)\n", + "6. Square-root transform" + ] }, { "cell_type": "markdown", "id": "load-10x-md", "metadata": {}, - "source": "### 3.1 Load 10X Data\n\n`sc.read_10x_mtx` reads each CellRanger output directory directly into an AnnData object. We label each sample with its time point, concatenate all samples, then compute QC metrics — including the fraction of mitochondrial gene counts needed for dead-cell removal later." + "source": [ + "### 3.1 Load 10X Data\n", + "\n", + "`sc.read_10x_mtx` reads each CellRanger output directory directly into an AnnData object. We label each sample with its time point, concatenate all samples, then compute QC metrics — including the fraction of mitochondrial gene counts needed for dead-cell removal later." + ] }, { "cell_type": "code", @@ -132,13 +213,44 @@ "id": "load-10x", "metadata": {}, "outputs": [], - "source": "samples = ['T0_1A', 'T2_3B', 'T4_5C', 'T6_7D', 'T8_9E']\nlabels = ['Day 00-03', 'Day 06-09', 'Day 12-15', 'Day 18-21', 'Day 24-27']\n\nadatas = []\nfor sample, label in zip(samples, labels):\n adata_sample = sc.read_10x_mtx(\n os.path.join(RAW_DATA_DIR, sample),\n var_names='gene_symbols',\n make_unique=True,\n cache=True,\n )\n adata_sample.obs['time_label'] = label\n adatas.append(adata_sample)\n\nadata = sc.concat(adatas, merge='same')\nadata.obs_names_make_unique()\n\n# Compute QC metrics (library size, mitochondrial gene fraction)\nadata.var['mt'] = adata.var_names.str.startswith('MT-')\nsc.pp.calculate_qc_metrics(adata, qc_vars=['mt'], percent_top=None, log1p=False, inplace=True)\n\nprint('Cells per sample:')\nfor label in labels:\n n = (adata.obs['time_label'] == label).sum()\n print(f' {label}: {n}')\nprint(f'\\nTotal: {adata.n_obs} cells × {adata.n_vars} genes')" + "source": [ + "samples = ['T0_1A', 'T2_3B', 'T4_5C', 'T6_7D', 'T8_9E']\n", + "labels = ['Day 00-03', 'Day 06-09', 'Day 12-15', 'Day 18-21', 'Day 24-27']\n", + "\n", + "adatas = []\n", + "for sample, label in zip(samples, labels):\n", + " adata_sample = sc.read_10x_mtx(\n", + " os.path.join(RAW_DATA_DIR, sample),\n", + " var_names='gene_symbols',\n", + " make_unique=True,\n", + " cache=True,\n", + " )\n", + " adata_sample.obs['time_label'] = label\n", + " adatas.append(adata_sample)\n", + "\n", + "adata = sc.concat(adatas, merge='same')\n", + "adata.obs_names_make_unique()\n", + "\n", + "# Compute QC metrics (library size, mitochondrial gene fraction)\n", + "adata.var['mt'] = adata.var_names.str.startswith('MT-')\n", + "sc.pp.calculate_qc_metrics(adata, qc_vars=['mt'], percent_top=None, log1p=False, inplace=True)\n", + "\n", + "print('Cells per sample:')\n", + "for label in labels:\n", + " n = (adata.obs['time_label'] == label).sum()\n", + " print(f' {label}: {n}')\n", + "print(f'\\nTotal: {adata.n_obs} cells × {adata.n_vars} genes')" + ] }, { "cell_type": "markdown", "id": "lib-size-filter-md", "metadata": {}, - "source": "### 3.2 Library Size Filtering\n\nWe remove cells in the top and bottom 20% of library sizes **within each sample**. This removes empty droplets and potential doublets while accounting for the fact that library size correlates with sample." + "source": [ + "### 3.2 Library Size Filtering\n", + "\n", + "We remove cells in the top and bottom 20% of library sizes **within each sample**. This removes empty droplets and potential doublets while accounting for the fact that library size correlates with sample." + ] }, { "cell_type": "code", @@ -146,7 +258,33 @@ "id": "lib-size-plot", "metadata": {}, "outputs": [], - "source": "min_percentile = 20\nmax_percentile = 80\n\nfig, axes = plt.subplots(2, 3, figsize=(15, 8))\naxes = axes.flatten()\n\nfor idx, label in enumerate(labels):\n sample_counts = adata.obs.loc[adata.obs['time_label'] == label, 'total_counts']\n t_min = np.percentile(sample_counts, min_percentile)\n t_max = np.percentile(sample_counts, max_percentile)\n\n axes[idx].hist(sample_counts, bins=50, alpha=0.7, edgecolor='black', log=True)\n axes[idx].axvline(t_min, color='red', linestyle='--', linewidth=2,\n label=f'{min_percentile}th: {t_min:.0f}')\n axes[idx].axvline(t_max, color='blue', linestyle='--', linewidth=2,\n label=f'{max_percentile}th: {t_max:.0f}')\n axes[idx].set_xlabel('Library Size (Total Counts)')\n axes[idx].set_ylabel('Number of Cells')\n axes[idx].set_title(label)\n axes[idx].legend(fontsize=8)\n axes[idx].grid(alpha=0.3)\n\naxes[-1].axis('off')\nplt.tight_layout()\nplt.show()" + "source": [ + "min_percentile = 20\n", + "max_percentile = 80\n", + "\n", + "fig, axes = plt.subplots(2, 3, figsize=(15, 8))\n", + "axes = axes.flatten()\n", + "\n", + "for idx, label in enumerate(labels):\n", + " sample_counts = adata.obs.loc[adata.obs['time_label'] == label, 'total_counts']\n", + " t_min = np.percentile(sample_counts, min_percentile)\n", + " t_max = np.percentile(sample_counts, max_percentile)\n", + "\n", + " axes[idx].hist(sample_counts, bins=50, alpha=0.7, edgecolor='black', log=True)\n", + " axes[idx].axvline(t_min, color='red', linestyle='--', linewidth=2,\n", + " label=f'{min_percentile}th: {t_min:.0f}')\n", + " axes[idx].axvline(t_max, color='blue', linestyle='--', linewidth=2,\n", + " label=f'{max_percentile}th: {t_max:.0f}')\n", + " axes[idx].set_xlabel('Library Size (Total Counts)')\n", + " axes[idx].set_ylabel('Number of Cells')\n", + " axes[idx].set_title(label)\n", + " axes[idx].legend(fontsize=8)\n", + " axes[idx].grid(alpha=0.3)\n", + "\n", + "axes[-1].axis('off')\n", + "plt.tight_layout()\n", + "plt.show()" + ] }, { "cell_type": "code", @@ -154,13 +292,31 @@ "id": "lib-size-filter", "metadata": {}, "outputs": [], - "source": "cells_to_keep = []\n\nfor label in labels:\n sample_counts = adata.obs.loc[adata.obs['time_label'] == label, 'total_counts']\n t_min = np.percentile(sample_counts, min_percentile)\n t_max = np.percentile(sample_counts, max_percentile)\n keep = sample_counts[(sample_counts >= t_min) & (sample_counts <= t_max)].index\n cells_to_keep.extend(keep.tolist())\n print(f'{label}: keeping {len(keep)}/{len(sample_counts)} cells '\n f'(range: {t_min:.0f}–{t_max:.0f})')\n\nadata = adata[cells_to_keep, :].copy()\nprint(f'\\nCells after library-size filtering: {adata.n_obs}')" + "source": [ + "cells_to_keep = []\n", + "\n", + "for label in labels:\n", + " sample_counts = adata.obs.loc[adata.obs['time_label'] == label, 'total_counts']\n", + " t_min = np.percentile(sample_counts, min_percentile)\n", + " t_max = np.percentile(sample_counts, max_percentile)\n", + " keep = sample_counts[(sample_counts >= t_min) & (sample_counts <= t_max)].index\n", + " cells_to_keep.extend(keep.tolist())\n", + " print(f'{label}: keeping {len(keep)}/{len(sample_counts)} cells '\n", + " f'(range: {t_min:.0f}–{t_max:.0f})')\n", + "\n", + "adata = adata[cells_to_keep, :].copy()\n", + "print(f'\\nCells after library-size filtering: {adata.n_obs}')" + ] }, { "cell_type": "markdown", "id": "rare-genes-md", "metadata": {}, - "source": "### 3.3 Remove Rare Genes\n\nGenes expressed in 10 or fewer cells are unlikely to be biologically informative and are removed." + "source": [ + "### 3.3 Remove Rare Genes\n", + "\n", + "Genes expressed in 10 or fewer cells are unlikely to be biologically informative and are removed." + ] }, { "cell_type": "code", @@ -168,13 +324,20 @@ "id": "rare-genes", "metadata": {}, "outputs": [], - "source": "sc.pp.filter_genes(adata, min_cells=10)\nprint(f'Genes after rare-gene filtering: {adata.n_vars}')" + "source": [ + "sc.pp.filter_genes(adata, min_cells=10)\n", + "print(f'Genes after rare-gene filtering: {adata.n_vars}')" + ] }, { "cell_type": "markdown", "id": "normalise-md", "metadata": {}, - "source": "### 3.4 Library-Size Normalisation\n\nDivide each cell by its total count and rescale by the median library size to make cells comparable." + "source": [ + "### 3.4 Library-Size Normalisation\n", + "\n", + "Divide each cell by its total count and rescale by the median library size to make cells comparable." + ] }, { "cell_type": "code", @@ -182,13 +345,20 @@ "id": "normalise", "metadata": {}, "outputs": [], - "source": "sc.pp.normalize_total(adata, target_sum=np.median(adata.obs['total_counts']))\nprint('Library-size normalisation complete.')" + "source": [ + "sc.pp.normalize_total(adata, target_sum=np.median(adata.obs['total_counts']))\n", + "print('Library-size normalisation complete.')" + ] }, { "cell_type": "markdown", "id": "dead-cells-md", "metadata": {}, - "source": "### 3.5 Remove Dead Cells\n\nDead cells show elevated mitochondrial RNA expression. We remove cells in the top 90th percentile of `pct_counts_mt`, which was already computed in step 3.1." + "source": [ + "### 3.5 Remove Dead Cells\n", + "\n", + "Dead cells show elevated mitochondrial RNA expression. We remove cells in the top 90th percentile of `pct_counts_mt`, which was already computed in step 3.1." + ] }, { "cell_type": "code", @@ -196,7 +366,25 @@ "id": "dead-cells-plot", "metadata": {}, "outputs": [], - "source": "mito_pct = adata.obs['pct_counts_mt']\nmito_percentile = 90\nmito_threshold = np.percentile(mito_pct, mito_percentile)\n\nprint(f'Mitochondrial genes found: {adata.var[\"mt\"].sum()}')\n\nfig, ax = plt.subplots(figsize=(8, 5))\nax.hist(mito_pct, bins=50, alpha=0.7, edgecolor='black')\nax.axvline(mito_threshold, color='red', linestyle='--', linewidth=2,\n label=f'{mito_percentile}th percentile: {mito_threshold:.1f}%')\nax.set_xlabel('Mitochondrial Gene %')\nax.set_ylabel('Number of Cells')\nax.set_title('Mitochondrial Content Distribution')\nax.legend()\nax.grid(alpha=0.3)\nplt.tight_layout()\nplt.show()" + "source": [ + "mito_pct = adata.obs['pct_counts_mt']\n", + "mito_percentile = 90\n", + "mito_threshold = np.percentile(mito_pct, mito_percentile)\n", + "\n", + "print(f'Mitochondrial genes found: {adata.var[\"mt\"].sum()}')\n", + "\n", + "fig, ax = plt.subplots(figsize=(8, 5))\n", + "ax.hist(mito_pct, bins=50, alpha=0.7, edgecolor='black')\n", + "ax.axvline(mito_threshold, color='red', linestyle='--', linewidth=2,\n", + " label=f'{mito_percentile}th percentile: {mito_threshold:.1f}%')\n", + "ax.set_xlabel('Mitochondrial Gene %')\n", + "ax.set_ylabel('Number of Cells')\n", + "ax.set_title('Mitochondrial Content Distribution')\n", + "ax.legend()\n", + "ax.grid(alpha=0.3)\n", + "plt.tight_layout()\n", + "plt.show()" + ] }, { "cell_type": "code", @@ -204,13 +392,20 @@ "id": "dead-cells-filter", "metadata": {}, "outputs": [], - "source": "adata = adata[adata.obs['pct_counts_mt'] < mito_threshold].copy()\nprint(f'Cells after dead-cell removal: {adata.n_obs}')" + "source": [ + "adata = adata[adata.obs['pct_counts_mt'] < mito_threshold].copy()\n", + "print(f'Cells after dead-cell removal: {adata.n_obs}')" + ] }, { "cell_type": "markdown", "id": "sqrt-md", "metadata": {}, - "source": "### 3.6 Square-Root Transform\n\nWe use the square-root transform instead of log1p to stabilise variance. It has the same compressive shape as log but is stable at zero." + "source": [ + "### 3.6 Square-Root Transform\n", + "\n", + "We use the square-root transform instead of log1p to stabilise variance. It has the same compressive shape as log but is stable at zero." + ] }, { "cell_type": "code", @@ -218,13 +413,21 @@ "id": "sqrt", "metadata": {}, "outputs": [], - "source": "import scipy.sparse as sp\nadata.X = np.sqrt(adata.X.toarray() if sp.issparse(adata.X) else adata.X)\nprint(f'Preprocessing complete. Final matrix: {adata.n_obs} cells × {adata.n_vars} genes')" + "source": [ + "import scipy.sparse as sp\n", + "adata.X = np.sqrt(adata.X.toarray() if sp.issparse(adata.X) else adata.X)\n", + "print(f'Preprocessing complete. Final matrix: {adata.n_obs} cells × {adata.n_vars} genes')" + ] }, { "cell_type": "markdown", "id": "anndata-md", "metadata": {}, - "source": "## 4. AnnData Summary\n\nAll preprocessing is complete. The AnnData object is ready for dimensionality reduction." + "source": [ + "## 4. AnnData Summary\n", + "\n", + "All preprocessing is complete. The AnnData object is ready for dimensionality reduction." + ] }, { "cell_type": "code", @@ -232,7 +435,9 @@ "id": "anndata", "metadata": {}, "outputs": [], - "source": "print(adata)" + "source": [ + "print(adata)" + ] }, { "cell_type": "markdown", @@ -289,18 +494,22 @@ "id": "phate-plot", "metadata": {}, "outputs": [], - "source": "sc.pl.embedding(\n adata, basis='phate', color='time_label',\n cmap='Spectral', title='Embryoid Body — PHATE embedding',\n)" + "source": [ + "sc.pl.embedding(\n", + " adata, basis='phate', color='time_label',\n", + " cmap='Spectral', title='Embryoid Body — PHATE embedding',\n", + ")" + ] }, { "cell_type": "markdown", "id": "mioflow-prep-md", "metadata": {}, "source": [ - "## 6. Prepare Data for MIOFlow\n", + "## 6. Train GAGA Autoencoder\n", "\n", - "MIOFlow expects a DataFrame with:\n", - "- One column per embedding dimension (`d1`, `d2`, …)\n", - "- A `samples` column with an **integer** time-bin label for each cell" + "MIOFlow 2.0 uses a Geometric Autoencoder (GAGA) to learn a mapping from the high-dimensional PCA space to a 2D geometric latent space that approximates the PHATE distances. \n", + "This allows MIOFlow to natively decode trajectories back to gene space." ] }, { @@ -310,19 +519,20 @@ "metadata": {}, "outputs": [], "source": [ - "# Create integer time bins: Day 00-03 → 0, Day 06-09 → 1, …, Day 24-27 → 4\n", - "adata.obs['discrete_time'], _ = pd.factorize(adata.obs['time_label'])\n", + "from mioflow.gaga import fit_gaga\n", "\n", - "# Build the input DataFrame (PHATE dims + samples column)\n", - "n_phate = adata.obsm['X_phate'].shape[1]\n", - "mioflow_df = pd.DataFrame(\n", - " adata.obsm['X_phate'],\n", - " columns=[f'd{i}' for i in range(1, n_phate + 1)],\n", - ")\n", - "mioflow_df['samples'] = adata.obs['discrete_time'].values\n", + "# Ensure time bins are created\n", + "adata.obs['discrete_time'], _ = pd.factorize(adata.obs['time_label'])\n", "\n", - "print(mioflow_df.head())\n", - "print(f'\\nTime bins: {sorted(mioflow_df[\"samples\"].unique())}')" + "print(\"Training GAGA Autoencoder...\")\n", + "gaga_model = fit_gaga(\n", + " X_pca=adata.obsm['X_pca'],\n", + " X_phate=adata.obsm['X_phate'],\n", + " latent_dim=2,\n", + " encoder_epochs=150,\n", + " decoder_epochs=150,\n", + " device='cuda' if use_cuda else 'cpu'\n", + ")" ] }, { @@ -334,8 +544,10 @@ "source": [ "fig, ax = plt.subplots(figsize=(8, 6))\n", "sns.scatterplot(\n", - " data=mioflow_df, x='d1', y='d2',\n", - " hue='samples', palette='viridis',\n", + " x=adata.obsm['X_phate'][:, 0], \n", + " y=adata.obsm['X_phate'][:, 1],\n", + " hue=adata.obs['time_label'], \n", + " palette='viridis',\n", " s=3, ax=ax,\n", ")\n", "ax.set_title('PHATE embedding coloured by discrete time bin')\n", @@ -369,13 +581,6 @@ "metadata": {}, "outputs": [], "source": [ - "# Model architecture\n", - "MODEL_CONFIG = {\n", - " 'layers': [16, 32, 16],\n", - " 'activation': 'CELU',\n", - " 'use_cuda': use_cuda,\n", - "}\n", - "\n", "# Training\n", "TRAINING_CONFIG = {\n", " 'n_epochs': 40,\n", @@ -419,10 +624,12 @@ "source": [ "mioflow_operator = MIOFlow(\n", " adata,\n", - " input_df=mioflow_df,\n", - " obsm_key='X_phate',\n", + " gaga_model=gaga_model,\n", + " gaga_input_key='X_pca',\n", + " obs_time_key='discrete_time',\n", " debug_level='info',\n", - " model_config=MODEL_CONFIG,\n", + " use_cuda=use_cuda,\n", + " hidden_dim=32,\n", " **TRAINING_CONFIG,\n", " **OPTIMIZATION_CONFIG,\n", " **DATA_CONFIG,\n", @@ -467,11 +674,36 @@ "metadata": {}, "outputs": [], "source": [ - "plot_losses(\n", - " mioflow.local_losses,\n", - " mioflow.batch_losses,\n", - " mioflow.globe_losses,\n", - ")" + "fig, axes = plt.subplots(1, 2, figsize=(14, 4))\n", + "\n", + "if hasattr(mioflow, 'losses') and mioflow.losses.get('total_loss'):\n", + " epochs = mioflow.losses['epoch']\n", + "\n", + " # Left: Total Loss\n", + " axes[0].plot(epochs, mioflow.losses['total_loss'], label='Total Loss', color='black')\n", + " axes[0].set_title('Global Training Loss')\n", + " axes[0].set_xlabel('Epoch')\n", + " axes[0].set_ylabel('Loss (Log Scale)')\n", + " axes[0].set_yscale('log')\n", + " axes[0].grid(True, alpha=0.3)\n", + " axes[0].legend()\n", + "\n", + " # Right: Loss Components\n", + " if 'ot_loss' in mioflow.losses:\n", + " axes[1].plot(epochs, mioflow.losses['ot_loss'], label='Optimal Transport', color='tab:orange')\n", + " if 'density_loss' in mioflow.losses:\n", + " axes[1].plot(epochs, mioflow.losses['density_loss'], label='Density', color='tab:green')\n", + " if 'energy_loss' in mioflow.losses:\n", + " axes[1].plot(epochs, mioflow.losses['energy_loss'], label='Energy', color='tab:red')\n", + " \n", + " axes[1].set_title('Loss Components')\n", + " axes[1].set_xlabel('Epoch')\n", + " axes[1].set_yscale('log')\n", + " axes[1].grid(True, alpha=0.3)\n", + " axes[1].legend()\n", + "\n", + "plt.tight_layout()\n", + "plt.show()" ] }, { @@ -481,7 +713,7 @@ "source": [ "## 11. Visualise Trajectories\n", "\n", - "`mioflow.trajectories` has shape `(n_bins, n_trajectories, n_dims)` in normalised PHATE space. We denormalise before plotting." + "MIOFlow natively outputs trajectories in the GAGA latent space. We visualize them here by plotting the background cells in the GAGA latent space." ] }, { @@ -501,15 +733,23 @@ "metadata": {}, "outputs": [], "source": [ - "# Denormalise trajectories and original data back to PHATE scale\n", - "traj_pts = mioflow.trajectories * mioflow.std_vals + mioflow.mean_vals\n", - "dim_cols = [c for c in mioflow.df.columns if c != 'samples']\n", - "true_data = mioflow.df[dim_cols].values * mioflow.std_vals + mioflow.mean_vals\n", + "traj_pts = mioflow.trajectories\n", + "\n", + "# Generate GAGA latent embeddings for background cells\n", + "gaga_model.eval()\n", + "import torch\n", + "scaler = getattr(gaga_model, 'input_scaler', None)\n", + "X_raw = adata.obsm['X_pca'].astype(np.float32)\n", + "X_scaled = scaler.transform(X_raw) if scaler is not None else X_raw\n", + "with torch.no_grad():\n", + " gaga_embedding = gaga_model.encode(torch.tensor(X_scaled, device='cuda' if use_cuda else 'cpu')).cpu().numpy()\n", + "\n", + "time_bins = adata.obs['discrete_time'].values\n", "\n", "fig, ax = plt.subplots(figsize=(10, 8))\n", "sc = ax.scatter(\n", - " true_data[:, 0], true_data[:, 1],\n", - " c=mioflow.df['samples'].values, cmap='viridis', s=1, alpha=0.5,\n", + " gaga_embedding[:, 0], gaga_embedding[:, 1],\n", + " c=time_bins, cmap='viridis', s=1, alpha=0.5,\n", ")\n", "plt.colorbar(sc, ax=ax, label='time bin')\n", "\n", @@ -520,9 +760,9 @@ " arrowprops=dict(arrowstyle='->', color='black', lw=0.5, mutation_scale=10),\n", " )\n", "\n", - "ax.set_title('MIOFlow trajectories on PHATE embedding')\n", - "ax.set_xlabel('PHATE 1')\n", - "ax.set_ylabel('PHATE 2')\n", + "ax.set_title('MIOFlow trajectories on GAGA latent embedding')\n", + "ax.set_xlabel('GAGA 1')\n", + "ax.set_ylabel('GAGA 2')\n", "plt.tight_layout()\n", "plt.show()" ] @@ -569,6 +809,7 @@ "metadata": {}, "outputs": [], "source": [ + "import scanpy as sc\n", "sc.pp.highly_variable_genes(adata, n_top_genes=25)\n", "example_genes = adata.var_names[adata.var['highly_variable']]\n", "example_gene_mask = adata.var_names.isin(example_genes)\n", @@ -695,8 +936,8 @@ "metadata": {}, "outputs": [], "source": [ - "target_x = -0.006 # PHATE 1 coordinate of the desired trajectory endpoint\n", - "target_y = 0.020 # PHATE 2 coordinate of the desired trajectory endpoint\n", + "target_x = -0.006 # Adjust these coordinates based on the new GAGA plot!\n", + "target_y = 0.020 \n", "\n", "distances = [\n", " np.sqrt((traj[-1, 0] - target_x) ** 2 + (traj[-1, 1] - target_y) ** 2)\n", @@ -706,7 +947,7 @@ "print(f'Selected trajectory #{highlight_idx} (distance to target: {distances[highlight_idx]:.4f})')\n", "\n", "fig, ax = plt.subplots(figsize=(10, 8))\n", - "ax.scatter(true_data[:, 0], true_data[:, 1], c=mioflow.df['samples'].values, cmap='viridis', s=1, alpha=0.5)\n", + "ax.scatter(gaga_embedding[:, 0], gaga_embedding[:, 1], c=time_bins, cmap='viridis', s=1, alpha=0.5)\n", "\n", "for i, traj in enumerate(np.transpose(traj_pts, axes=(1, 0, 2))):\n", " colour = 'red' if i == highlight_idx else 'black'\n", @@ -756,5 +997,29 @@ "plt.show()" ] } - ] -} \ No newline at end of file + ], + "metadata": { + "colab": { + "provenance": [] + }, + "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", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.25" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}