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 +}