diff --git a/.gitignore b/.gitignore
index c63ff4b..d4fc55d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,5 +18,8 @@ auto_ver.sh
*.yaml
*.png
*.jpg
+# whitelist images committed to the docs static dir
+!docs/source/_static/img/*.png
+!docs/source/_static/img/*.jpg
.venv*
.nfs*
diff --git a/README.md b/README.md
index df12069..bd01045 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,7 @@
+ 
@@ -38,26 +39,23 @@ pip install scmidas
## Quick start
-The MIDAS workflow is four calls. The snippet below is an API sketch — replace `...` with your data and refer to the [tutorials](https://scmidas.readthedocs.io/en/latest/) for runnable end-to-end examples.
+A bundled 1600-cell PBMC RNA+ADT mosaic dataset lets you go from `pip install` to a UMAP in about a minute on a single GPU — no extra downloads, no config files. Click the [Colab badge](https://colab.research.google.com/github/labomics/midas/blob/main/examples/quickstart.ipynb) to run it without installing anything, or copy the snippet:
```python
-from scmidas.config import load_config
-from scmidas.model import MIDAS
+import scmidas
-# 1. Build a model bound to a mosaic dataset.
-# Input is either a directory of per-batch MTX matrices, or a MuData
-# object via MIDAS.configure_data_from_mdata(...).
-model = MIDAS.configure_data_from_dir(load_config(), ..., transform={'atac': 'binarize'})
+mdata = scmidas.datasets.quickstart() # bundled toy MuData
+model = scmidas.integrate(mdata) # ~1 min on a mid-range GPU
+out = model.predict(joint_latent=True) # latent embeddings per batch
+```
-# 2. Train.
-model.train(max_epochs=2000)
+This produces lineage-separated clusters that mix freely across batches:
-# 3. Predict — latent embeddings (z_c, z_u) and imputed counts per batch.
-pred = model.predict()
+
+
+
MuData object with n_obs × n_vars = 1600 × 724\n",
+ " obs:\t'batch', 'celltype'\n",
+ " 2 modalities\n",
+ " rna:\t1200 × 500\n",
+ " obs:\t'batch', 'celltype'\n",
+ " adt:\t1200 × 224\n",
+ " obs:\t'batch', 'celltype'
"
+ ],
+ "text/plain": [
+ "MuData object with n_obs × n_vars = 1600 × 724\n",
+ " obs:\t'batch', 'celltype'\n",
+ " 2 modalities\n",
+ " rna:\t1200 × 500\n",
+ " obs:\t'batch', 'celltype'\n",
+ " adt:\t1200 × 224\n",
+ " obs:\t'batch', 'celltype'"
+ ]
+ },
+ "execution_count": 1,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "import warnings\n",
+ "warnings.filterwarnings('ignore')\n",
+ "import logging\n",
+ "logging.basicConfig(level=logging.INFO)\n",
+ "\n",
+ "import scmidas\n",
+ "\n",
+ "mdata = scmidas.datasets.quickstart()\n",
+ "mdata"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "602dc63b",
+ "metadata": {},
+ "source": [
+ "## 3. Train\n",
+ "\n",
+ "One call. `scmidas.integrate(...)` infers feature dimensions from the MuData and uses defaults that produce a clean UMAP for this toy dataset in about a minute."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "e91fadfd",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-05-03T05:44:25.940862Z",
+ "iopub.status.busy": "2026-05-03T05:44:25.940266Z",
+ "iopub.status.idle": "2026-05-03T05:45:29.365469Z",
+ "shell.execute_reply": "2026-05-03T05:45:29.364139Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "INFO:scmidas.config:The model is initialized with the default configurations.\n",
+ "INFO:scmidas.api:scmidas.integrate(): toy-tuned defaults — batch_size=128, max_epochs=65, lr=0.0003. For real datasets, override max_epochs (e.g. 2000) and consider batch_size=256.\n",
+ "INFO:scmidas.model:Input data: \n",
+ " #CELL #RNA #ADT\n",
+ "p1_0 400 500.0 NaN\n",
+ "p2_0 400 NaN 224.0\n",
+ "p3_0 400 500.0 224.0\n",
+ "p4_0 400 500.0 224.0\n",
+ "INFO: 💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
+ "INFO:lightning.pytorch.utilities.rank_zero:💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
+ "INFO: GPU available: True (cuda), used: True\n",
+ "INFO:lightning.pytorch.utilities.rank_zero:GPU available: True (cuda), used: True\n",
+ "INFO: TPU available: False, using: 0 TPU cores\n",
+ "INFO:lightning.pytorch.utilities.rank_zero:TPU available: False, using: 0 TPU cores\n",
+ "INFO: HPU available: False, using: 0 HPUs\n",
+ "INFO:lightning.pytorch.utilities.rank_zero:HPU available: False, using: 0 HPUs\n",
+ "INFO: LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
+ "INFO:lightning.pytorch.accelerators.cuda:LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
+ "INFO: \n",
+ " | Name | Type | Params | Mode \n",
+ "-----------------------------------------------\n",
+ "0 | net | VAE | 1.8 M | train\n",
+ "1 | dsc | Discriminator | 39.2 K | train\n",
+ "-----------------------------------------------\n",
+ "1.8 M Trainable params\n",
+ "0 Non-trainable params\n",
+ "1.8 M Total params\n",
+ "7.240 Total estimated model params size (MB)\n",
+ "144 Modules in train mode\n",
+ "0 Modules in eval mode\n",
+ "INFO:lightning.pytorch.callbacks.model_summary:\n",
+ " | Name | Type | Params | Mode \n",
+ "-----------------------------------------------\n",
+ "0 | net | VAE | 1.8 M | train\n",
+ "1 | dsc | Discriminator | 39.2 K | train\n",
+ "-----------------------------------------------\n",
+ "1.8 M Trainable params\n",
+ "0 Non-trainable params\n",
+ "1.8 M Total params\n",
+ "7.240 Total estimated model params size (MB)\n",
+ "144 Modules in train mode\n",
+ "0 Modules in eval mode\n",
+ "INFO:scmidas.model:Total number of samples: 1600 from 4 datasets.\n",
+ "INFO:scmidas.model:Using MultiBatchSampler for data loading.\n",
+ "INFO:scmidas.model:DataLoader created with batch size 128 and 20 workers.\n"
+ ]
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "30186375b1a0421e92b5cf6e4c5dfda4",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Training: | | 0/? [00:00"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "import anndata as ad\n",
+ "import matplotlib.pyplot as plt\n",
+ "import scanpy as sc\n",
+ "\n",
+ "ad_z = ad.AnnData(X=z_c)\n",
+ "ad_z.obs_names = all_ids\n",
+ "batches, celltypes = [], []\n",
+ "for c in all_ids:\n",
+ " if c in mdata['rna'].obs_names:\n",
+ " batches.append(mdata['rna'].obs.loc[c, 'batch'])\n",
+ " celltypes.append(mdata['rna'].obs.loc[c, 'celltype'])\n",
+ " else:\n",
+ " batches.append(mdata['adt'].obs.loc[c, 'batch'])\n",
+ " celltypes.append(mdata['adt'].obs.loc[c, 'celltype'])\n",
+ "ad_z.obs['batch'] = batches\n",
+ "ad_z.obs['celltype'] = celltypes\n",
+ "\n",
+ "sc.pp.neighbors(ad_z, n_neighbors=30)\n",
+ "sc.tl.umap(ad_z, random_state=42, min_dist=0.3)\n",
+ "\n",
+ "fig, axes = plt.subplots(1, 2, figsize=(13, 5.5))\n",
+ "sc.pl.umap(ad_z, color='celltype', ax=axes[0], show=False, frameon=False, title='Cell type')\n",
+ "sc.pl.umap(ad_z, color='batch', ax=axes[1], show=False, frameon=False, title='Batch')\n",
+ "plt.tight_layout()\n",
+ "plt.show()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "be94ec90",
+ "metadata": {},
+ "source": [
+ "## Next steps\n",
+ "\n",
+ "- See `docs/source/tutorials/basics/` for the full demos (with real-sized data, multi-GPU training, prediction tasks like imputation, batch correction, modality translation).\n",
+ "- For your own MuData: `scmidas.integrate(your_mdata, max_epochs=2000, batch_size=256)` (paper-default hyperparameters)."
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.12.11"
+ },
+ "widgets": {
+ "application/vnd.jupyter.widget-state+json": {
+ "state": {
+ "30186375b1a0421e92b5cf6e4c5dfda4": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_module_version": "2.0.0",
+ "model_name": "HBoxModel",
+ "state": {
+ "_dom_classes": [],
+ "_model_module": "@jupyter-widgets/controls",
+ "_model_module_version": "2.0.0",
+ "_model_name": "HBoxModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/controls",
+ "_view_module_version": "2.0.0",
+ "_view_name": "HBoxView",
+ "box_style": "",
+ "children": [
+ "IPY_MODEL_dd59645be4d9462ea87eeb3280ee9463",
+ "IPY_MODEL_c04fd7c0d62f4db1ac61bce7e3fc8987",
+ "IPY_MODEL_90c165cf45a04c5292db0b7eea9db213"
+ ],
+ "layout": "IPY_MODEL_4ee8a0e03b0c4957a25a128c2d50f5d4",
+ "tabbable": null,
+ "tooltip": null
+ }
+ },
+ "39ac3c9b2c9248e58188bde88e5d542a": {
+ "model_module": "@jupyter-widgets/base",
+ "model_module_version": "2.0.0",
+ "model_name": "LayoutModel",
+ "state": {
+ "_model_module": "@jupyter-widgets/base",
+ "_model_module_version": "2.0.0",
+ "_model_name": "LayoutModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/base",
+ "_view_module_version": "2.0.0",
+ "_view_name": "LayoutView",
+ "align_content": null,
+ "align_items": null,
+ "align_self": null,
+ "border_bottom": null,
+ "border_left": null,
+ "border_right": null,
+ "border_top": null,
+ "bottom": null,
+ "display": null,
+ "flex": "2",
+ "flex_flow": null,
+ "grid_area": null,
+ "grid_auto_columns": null,
+ "grid_auto_flow": null,
+ "grid_auto_rows": null,
+ "grid_column": null,
+ "grid_gap": null,
+ "grid_row": null,
+ "grid_template_areas": null,
+ "grid_template_columns": null,
+ "grid_template_rows": null,
+ "height": null,
+ "justify_content": null,
+ "justify_items": null,
+ "left": null,
+ "margin": null,
+ "max_height": null,
+ "max_width": null,
+ "min_height": null,
+ "min_width": null,
+ "object_fit": null,
+ "object_position": null,
+ "order": null,
+ "overflow": null,
+ "padding": null,
+ "right": null,
+ "top": null,
+ "visibility": null,
+ "width": null
+ }
+ },
+ "4ee8a0e03b0c4957a25a128c2d50f5d4": {
+ "model_module": "@jupyter-widgets/base",
+ "model_module_version": "2.0.0",
+ "model_name": "LayoutModel",
+ "state": {
+ "_model_module": "@jupyter-widgets/base",
+ "_model_module_version": "2.0.0",
+ "_model_name": "LayoutModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/base",
+ "_view_module_version": "2.0.0",
+ "_view_name": "LayoutView",
+ "align_content": null,
+ "align_items": null,
+ "align_self": null,
+ "border_bottom": null,
+ "border_left": null,
+ "border_right": null,
+ "border_top": null,
+ "bottom": null,
+ "display": "inline-flex",
+ "flex": null,
+ "flex_flow": "row wrap",
+ "grid_area": null,
+ "grid_auto_columns": null,
+ "grid_auto_flow": null,
+ "grid_auto_rows": null,
+ "grid_column": null,
+ "grid_gap": null,
+ "grid_row": null,
+ "grid_template_areas": null,
+ "grid_template_columns": null,
+ "grid_template_rows": null,
+ "height": null,
+ "justify_content": null,
+ "justify_items": null,
+ "left": null,
+ "margin": null,
+ "max_height": null,
+ "max_width": null,
+ "min_height": null,
+ "min_width": null,
+ "object_fit": null,
+ "object_position": null,
+ "order": null,
+ "overflow": null,
+ "padding": null,
+ "right": null,
+ "top": null,
+ "visibility": null,
+ "width": "100%"
+ }
+ },
+ "54721dddcd6445aab2ee8a974bc36b54": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_module_version": "2.0.0",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_model_module": "@jupyter-widgets/controls",
+ "_model_module_version": "2.0.0",
+ "_model_name": "ProgressStyleModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/base",
+ "_view_module_version": "2.0.0",
+ "_view_name": "StyleView",
+ "bar_color": null,
+ "description_width": ""
+ }
+ },
+ "737974cb592f48dd96dc726efe5b61c6": {
+ "model_module": "@jupyter-widgets/base",
+ "model_module_version": "2.0.0",
+ "model_name": "LayoutModel",
+ "state": {
+ "_model_module": "@jupyter-widgets/base",
+ "_model_module_version": "2.0.0",
+ "_model_name": "LayoutModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/base",
+ "_view_module_version": "2.0.0",
+ "_view_name": "LayoutView",
+ "align_content": null,
+ "align_items": null,
+ "align_self": null,
+ "border_bottom": null,
+ "border_left": null,
+ "border_right": null,
+ "border_top": null,
+ "bottom": null,
+ "display": null,
+ "flex": null,
+ "flex_flow": null,
+ "grid_area": null,
+ "grid_auto_columns": null,
+ "grid_auto_flow": null,
+ "grid_auto_rows": null,
+ "grid_column": null,
+ "grid_gap": null,
+ "grid_row": null,
+ "grid_template_areas": null,
+ "grid_template_columns": null,
+ "grid_template_rows": null,
+ "height": null,
+ "justify_content": null,
+ "justify_items": null,
+ "left": null,
+ "margin": null,
+ "max_height": null,
+ "max_width": null,
+ "min_height": null,
+ "min_width": null,
+ "object_fit": null,
+ "object_position": null,
+ "order": null,
+ "overflow": null,
+ "padding": null,
+ "right": null,
+ "top": null,
+ "visibility": null,
+ "width": null
+ }
+ },
+ "90c165cf45a04c5292db0b7eea9db213": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_module_version": "2.0.0",
+ "model_name": "HTMLModel",
+ "state": {
+ "_dom_classes": [],
+ "_model_module": "@jupyter-widgets/controls",
+ "_model_module_version": "2.0.0",
+ "_model_name": "HTMLModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/controls",
+ "_view_module_version": "2.0.0",
+ "_view_name": "HTMLView",
+ "description": "",
+ "description_allow_html": false,
+ "layout": "IPY_MODEL_e288a669839d41f39581cb43cdc38b85",
+ "placeholder": "",
+ "style": "IPY_MODEL_bbab4f6701f74e20b55d21f3119e0bc2",
+ "tabbable": null,
+ "tooltip": null,
+ "value": " 16/16 [00:00<00:00, 17.14it/s, v_num=0, loss_/recon_loss_step=414.0, loss_/kld_loss_step=61.10, loss_/consistency_loss_step=0.000, loss/net_step=411.0, loss/dsc_step=63.20, loss_/recon_loss_epoch=1.51e+3, loss_/kld_loss_epoch=79.20, loss_/consistency_loss_epoch=58.20, loss/net_epoch=1.57e+3, loss/dsc_epoch=78.00]"
+ }
+ },
+ "bbab4f6701f74e20b55d21f3119e0bc2": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_module_version": "2.0.0",
+ "model_name": "HTMLStyleModel",
+ "state": {
+ "_model_module": "@jupyter-widgets/controls",
+ "_model_module_version": "2.0.0",
+ "_model_name": "HTMLStyleModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/base",
+ "_view_module_version": "2.0.0",
+ "_view_name": "StyleView",
+ "background": null,
+ "description_width": "",
+ "font_size": null,
+ "text_color": null
+ }
+ },
+ "c04fd7c0d62f4db1ac61bce7e3fc8987": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_module_version": "2.0.0",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_dom_classes": [],
+ "_model_module": "@jupyter-widgets/controls",
+ "_model_module_version": "2.0.0",
+ "_model_name": "FloatProgressModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/controls",
+ "_view_module_version": "2.0.0",
+ "_view_name": "ProgressView",
+ "bar_style": "success",
+ "description": "",
+ "description_allow_html": false,
+ "layout": "IPY_MODEL_39ac3c9b2c9248e58188bde88e5d542a",
+ "max": 16.0,
+ "min": 0.0,
+ "orientation": "horizontal",
+ "style": "IPY_MODEL_54721dddcd6445aab2ee8a974bc36b54",
+ "tabbable": null,
+ "tooltip": null,
+ "value": 16.0
+ }
+ },
+ "c2d1dd3030f7426bb2485c04c3d7a500": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_module_version": "2.0.0",
+ "model_name": "HTMLStyleModel",
+ "state": {
+ "_model_module": "@jupyter-widgets/controls",
+ "_model_module_version": "2.0.0",
+ "_model_name": "HTMLStyleModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/base",
+ "_view_module_version": "2.0.0",
+ "_view_name": "StyleView",
+ "background": null,
+ "description_width": "",
+ "font_size": null,
+ "text_color": null
+ }
+ },
+ "dd59645be4d9462ea87eeb3280ee9463": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_module_version": "2.0.0",
+ "model_name": "HTMLModel",
+ "state": {
+ "_dom_classes": [],
+ "_model_module": "@jupyter-widgets/controls",
+ "_model_module_version": "2.0.0",
+ "_model_name": "HTMLModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/controls",
+ "_view_module_version": "2.0.0",
+ "_view_name": "HTMLView",
+ "description": "",
+ "description_allow_html": false,
+ "layout": "IPY_MODEL_737974cb592f48dd96dc726efe5b61c6",
+ "placeholder": "",
+ "style": "IPY_MODEL_c2d1dd3030f7426bb2485c04c3d7a500",
+ "tabbable": null,
+ "tooltip": null,
+ "value": "Epoch 64: 100%"
+ }
+ },
+ "e288a669839d41f39581cb43cdc38b85": {
+ "model_module": "@jupyter-widgets/base",
+ "model_module_version": "2.0.0",
+ "model_name": "LayoutModel",
+ "state": {
+ "_model_module": "@jupyter-widgets/base",
+ "_model_module_version": "2.0.0",
+ "_model_name": "LayoutModel",
+ "_view_count": null,
+ "_view_module": "@jupyter-widgets/base",
+ "_view_module_version": "2.0.0",
+ "_view_name": "LayoutView",
+ "align_content": null,
+ "align_items": null,
+ "align_self": null,
+ "border_bottom": null,
+ "border_left": null,
+ "border_right": null,
+ "border_top": null,
+ "bottom": null,
+ "display": null,
+ "flex": null,
+ "flex_flow": null,
+ "grid_area": null,
+ "grid_auto_columns": null,
+ "grid_auto_flow": null,
+ "grid_auto_rows": null,
+ "grid_column": null,
+ "grid_gap": null,
+ "grid_row": null,
+ "grid_template_areas": null,
+ "grid_template_columns": null,
+ "grid_template_rows": null,
+ "height": null,
+ "justify_content": null,
+ "justify_items": null,
+ "left": null,
+ "margin": null,
+ "max_height": null,
+ "max_width": null,
+ "min_height": null,
+ "min_width": null,
+ "object_fit": null,
+ "object_position": null,
+ "order": null,
+ "overflow": null,
+ "padding": null,
+ "right": null,
+ "top": null,
+ "visibility": null,
+ "width": null
+ }
+ }
+ },
+ "version_major": 2,
+ "version_minor": 0
+ }
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/pyproject.toml b/pyproject.toml
index 5d3e54c..ebcf800 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "scmidas"
-version = '0.1.19'
+version = '0.2.0'
description = "A torch-based integration method for single-cell multi-omic data."
readme = "README.md"
requires-python = ">=3.10"
@@ -101,7 +101,7 @@ include = ["scmidas*"]
# Ship the PEP 561 type-info marker so type checkers (mypy / pyright)
# treat scmidas as a typed package.
[tool.setuptools.package-data]
-scmidas = ["py.typed"]
+scmidas = ["py.typed", "data/*.h5mu"]
# --- Ruff (optional) ---
[tool.ruff]
diff --git a/scripts/build_quickstart_demo.py b/scripts/build_quickstart_demo.py
new file mode 100644
index 0000000..0280e33
--- /dev/null
+++ b/scripts/build_quickstart_demo.py
@@ -0,0 +1,169 @@
+"""Build the lightweight quickstart MuData from wnn_mosaic_8batch_mtx.
+
+Output: a single .h5mu file roughly 3-5 MB containing a representative
+mosaic subset (4 batches, ~600 cells each, 500 HVGs + all 224 ADT
+features). Reproducible (seed=42).
+"""
+from __future__ import annotations
+
+import argparse
+import warnings
+from pathlib import Path
+
+import anndata as ad
+import mudata as mu
+import numpy as np
+import pandas as pd
+import scanpy as sc
+from scipy.io import mmread
+
+warnings.filterwarnings('ignore')
+sc.settings.verbosity = 0
+
+SEED = 42
+N_CELLS_PER_BATCH = 400
+N_HVG = 500
+
+# Picked to preserve the full 1+1+2 mosaic structure of the source dataset
+# while keeping the result < 10MB. p3_0 brings NK over-representation;
+# p4_0 brings the more typical CD4/Mono/CD8 mix.
+SELECTED_BATCHES: list[str] = ['p1_0', 'p2_0', 'p3_0', 'p4_0']
+MODALITIES_PER_BATCH: dict[str, list[str]] = {
+ 'p1_0': ['rna'],
+ 'p2_0': ['adt'],
+ 'p3_0': ['rna', 'adt'],
+ 'p4_0': ['rna', 'adt'],
+}
+
+
+def stratified_subsample(celltypes: np.ndarray, n_target: int, seed: int) -> np.ndarray:
+ rng = np.random.default_rng(seed)
+ df = pd.DataFrame({'idx': np.arange(len(celltypes)), 'ct': celltypes})
+ n = len(celltypes)
+ if n_target >= n:
+ return df['idx'].values
+ chosen = []
+ for ct, group in df.groupby('ct'):
+ k = max(1, int(round(n_target * len(group) / n)))
+ k = min(k, len(group))
+ chosen.extend(rng.choice(group['idx'].values, size=k, replace=False).tolist())
+ chosen = np.array(chosen)
+ if len(chosen) > n_target:
+ chosen = rng.choice(chosen, size=n_target, replace=False)
+ elif len(chosen) < n_target:
+ leftover = np.setdiff1d(df['idx'].values, chosen)
+ extra = rng.choice(leftover, size=n_target - len(chosen), replace=False)
+ chosen = np.concatenate([chosen, extra])
+ return np.sort(chosen)
+
+
+def select_hvgs(src: Path, rna_names: list[str], rna_batches: list[str]) -> list[str]:
+ """Compute Seurat batch-aware HVGs across the *selected* RNA-having batches."""
+ adatas = []
+ for b in rna_batches:
+ m = mmread(src / f'data/{b}/mat/rna.mtx').tocsr().astype(np.float32)
+ a = ad.AnnData(X=m)
+ a.var_names = rna_names
+ a.obs_names = [f'{b}_{i}' for i in range(m.shape[0])]
+ a.obs['batch'] = b
+ adatas.append(a)
+ combined = ad.concat(adatas, join='outer')
+ combined.obs_names_make_unique()
+ sc.pp.normalize_total(combined, target_sum=1e4)
+ sc.pp.log1p(combined)
+ sc.pp.highly_variable_genes(combined, n_top_genes=N_HVG, batch_key='batch', flavor='seurat')
+ return combined.var.index[combined.var['highly_variable']].tolist()
+
+
+def build(src_dir: str, out_path: str) -> mu.MuData:
+ src = Path(src_dir)
+
+ rna_names = pd.read_csv(src / 'data/feat/feat_names_rna.csv', index_col=0)['x'].tolist()
+ adt_names = pd.read_csv(src / 'data/feat/feat_names_adt.csv', index_col=0)['x'].tolist()
+ print(f'Source: {len(rna_names)} RNA features, {len(adt_names)} ADT features')
+
+ rna_batches = [b for b in SELECTED_BATCHES if 'rna' in MODALITIES_PER_BATCH[b]]
+ adt_batches = [b for b in SELECTED_BATCHES if 'adt' in MODALITIES_PER_BATCH[b]]
+ print(f'Selected batches: {SELECTED_BATCHES}')
+ print(f' RNA batches: {rna_batches}')
+ print(f' ADT batches: {adt_batches}')
+
+ print(f'\nComputing {N_HVG} HVGs across selected RNA batches...')
+ hvgs = select_hvgs(src, rna_names, rna_batches)
+ print(f' selected {len(hvgs)} HVGs')
+
+ print(f'\nStratified subsample ({N_CELLS_PER_BATCH} cells/batch)...')
+ selected_idx: dict[str, np.ndarray] = {}
+ for b in SELECTED_BATCHES:
+ lbl = pd.read_csv(src / f'label/{b}.csv', index_col=0)
+ celltypes = lbl.iloc[:, 0].values
+ idx = stratified_subsample(celltypes, N_CELLS_PER_BATCH, SEED)
+ selected_idx[b] = idx
+ ct_dist = pd.Series(celltypes[idx]).value_counts().head(4).to_dict()
+ print(f' {b}: {len(idx)} cells | top: {ct_dist}')
+
+ hvg_mask = np.isin(rna_names, hvgs)
+
+ rna_adatas = []
+ adt_adatas = []
+ for b in SELECTED_BATCHES:
+ idx = selected_idx[b]
+ cell_ids = [f'{b}_{i}' for i in idx]
+ lbl = pd.read_csv(src / f'label/{b}.csv', index_col=0)
+ celltypes = lbl.iloc[:, 0].values[idx]
+
+ if 'rna' in MODALITIES_PER_BATCH[b]:
+ m = mmread(src / f'data/{b}/mat/rna.mtx').tocsr().astype(np.float32)
+ m = m[idx, :][:, hvg_mask]
+ a = ad.AnnData(X=m)
+ a.var_names = [n for n, k in zip(rna_names, hvg_mask) if k]
+ a.obs_names = cell_ids
+ a.obs['batch'] = b
+ a.obs['celltype'] = celltypes
+ rna_adatas.append(a)
+
+ if 'adt' in MODALITIES_PER_BATCH[b]:
+ m = mmread(src / f'data/{b}/mat/adt.mtx').tocsr().astype(np.float32)
+ m = m[idx, :]
+ a = ad.AnnData(X=m)
+ a.var_names = adt_names
+ a.obs_names = cell_ids
+ a.obs['batch'] = b
+ a.obs['celltype'] = celltypes
+ adt_adatas.append(a)
+
+ rna_full = ad.concat(rna_adatas, join='outer')
+ adt_full = ad.concat(adt_adatas, join='outer')
+ print(f'\nAnnData shapes: rna={rna_full.shape} adt={adt_full.shape}')
+
+ mdata = mu.MuData({'rna': rna_full, 'adt': adt_full})
+ mdata.update()
+
+ # Promote batch / celltype from modality obs to top-level mdata.obs so
+ # users only need to look in one place.
+ batch_per_cell: dict[str, str] = {}
+ ct_per_cell: dict[str, str] = {}
+ for mod_adata in [rna_full, adt_full]:
+ for cid, b, ct in zip(mod_adata.obs_names, mod_adata.obs['batch'], mod_adata.obs['celltype']):
+ batch_per_cell.setdefault(cid, b)
+ ct_per_cell.setdefault(cid, ct)
+ mdata.obs['batch'] = pd.Series(batch_per_cell)
+ mdata.obs['celltype'] = pd.Series(ct_per_cell)
+
+ out = Path(out_path)
+ out.parent.mkdir(parents=True, exist_ok=True)
+ mdata.write_h5mu(str(out))
+ size_mb = out.stat().st_size / 1024 / 1024
+ print(f'\nWrote {out}: {size_mb:.2f} MB')
+
+ return mdata
+
+
+if __name__ == '__main__':
+ p = argparse.ArgumentParser()
+ p.add_argument('--src', default='/tmp/midas-demo2/wnn_mosaic_8batch_mtx',
+ help='Path to extracted wnn_mosaic_8batch_mtx directory')
+ p.add_argument('--out', default='/tmp/quickstart_pbmc_mosaic.h5mu',
+ help='Output .h5mu path')
+ args = p.parse_args()
+ build(args.src, args.out)
diff --git a/src/scmidas/__init__.py b/src/scmidas/__init__.py
index e79d49e..4c96e8f 100644
--- a/src/scmidas/__init__.py
+++ b/src/scmidas/__init__.py
@@ -3,6 +3,8 @@
from .nn import *
from .utils import *
from .config import *
+from . import datasets
+from .api import integrate
from importlib.metadata import PackageNotFoundError, version as _pkg_version
try:
diff --git a/src/scmidas/api.py b/src/scmidas/api.py
new file mode 100644
index 0000000..8bc78a6
--- /dev/null
+++ b/src/scmidas/api.py
@@ -0,0 +1,127 @@
+"""High-level conveniences on top of the MIDAS class."""
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+import mudata as mu
+
+from .config import load_config
+from .model import MIDAS
+
+logger = logging.getLogger(__name__)
+
+
+# Tuned for the bundled quickstart dataset (PBMC mosaic, 1200 cells, 500
+# HVGs + 224 ADT). On a single mid-range GPU these defaults converge in
+# roughly one minute and produce a clean lineage-separated UMAP. They
+# are NOT general-purpose — for full datasets, fall back to the
+# paper defaults via ``MIDAS.configure_data_from_mdata(...).train()``.
+_QUICKSTART_DEFAULTS: dict[str, Any] = {
+ 'batch_size': 128,
+ 'max_epochs': 65,
+ 'lr_net': 3e-4,
+ 'lr_dsc': 3e-4,
+}
+
+
+def integrate(
+ mdata: mu.MuData,
+ *,
+ batch_key: str = 'batch',
+ max_epochs: Optional[int] = None,
+ batch_size: Optional[int] = None,
+ accelerator: str = 'auto',
+ devices: Any = 1,
+ strategy: str = 'auto',
+ save_model_path: str = './saved_models/scmidas',
+ seed: Optional[int] = 42,
+ **kwargs: Any,
+) -> MIDAS:
+ """One-call MIDAS pipeline for users who want a sensible default.
+
+ Internally this is just::
+
+ configs = load_config()
+ configs.update(quickstart_defaults)
+ MIDAS.configure_data_from_mdata(...).train(...)
+
+ so the surface and behaviour are identical to the longhand
+ pipeline; ``integrate`` simply fills in the easy-to-get-wrong
+ parameters with values that we know work on the bundled
+ ``scmidas.datasets.quickstart()`` data.
+
+ .. warning::
+ The default training hyperparameters (``batch_size=128``,
+ ``max_epochs=65``, ``lr=3e-4``) are tuned for the **toy
+ quickstart dataset** (1600 cells). They are **not appropriate
+ for real analyses** — for full datasets pass your own
+ ``max_epochs`` (typically 1000-2000) and consider letting
+ ``batch_size`` default back to 256.
+
+ Parameters:
+ mdata : MuData
+ Multi-modal single-cell data. Must have a top-level
+ ``mdata.obs[batch_key]`` column identifying batches.
+ batch_key : str
+ Column in ``mdata.obs`` that identifies the source batch.
+ max_epochs : int, optional
+ Training epochs. Default 65 (quickstart-tuned). For real
+ data, override with 1000-2000.
+ batch_size : int, optional
+ Mini-batch size. Default 128 (quickstart-tuned). For real
+ data, 256 is a more typical choice.
+ accelerator, devices, strategy
+ Forwarded to ``lightning.Trainer``. Default ``'auto'``
+ picks GPU if available.
+ save_model_path : str
+ Where to write checkpoints during training.
+ seed : int, optional
+ If not None, calls ``lightning.seed_everything(seed)``
+ before configuring data, so the run is reproducible.
+ **kwargs
+ Additional keyword arguments forwarded to
+ ``MIDAS.configure_data_from_mdata``.
+
+ Returns:
+ MIDAS:
+ A trained MIDAS model. Call ``.predict(...)`` to obtain
+ latent embeddings and/or imputed counts.
+ """
+ if seed is not None:
+ import lightning as L
+ L.seed_everything(seed, verbose=False)
+
+ configs = load_config()
+ configs['lr_net'] = _QUICKSTART_DEFAULTS['lr_net']
+ configs['lr_dsc'] = _QUICKSTART_DEFAULTS['lr_dsc']
+
+ bsz = batch_size if batch_size is not None else _QUICKSTART_DEFAULTS['batch_size']
+ eps = max_epochs if max_epochs is not None else _QUICKSTART_DEFAULTS['max_epochs']
+
+ dims_x = {m: [mdata[m].n_vars] for m in mdata.mod}
+
+ logger.info(
+ 'scmidas.integrate(): toy-tuned defaults — '
+ 'batch_size=%d, max_epochs=%d, lr=%g. '
+ 'For real datasets, override max_epochs (e.g. 2000) '
+ 'and consider batch_size=256.',
+ bsz, eps, _QUICKSTART_DEFAULTS['lr_net'],
+ )
+
+ model = MIDAS.configure_data_from_mdata(
+ configs=configs,
+ mdata=mdata,
+ dims_x=dims_x,
+ batch_key=batch_key,
+ batch_size=bsz,
+ save_model_path=save_model_path,
+ **kwargs,
+ )
+ model.train(
+ max_epochs=eps,
+ accelerator=accelerator,
+ devices=devices,
+ strategy=strategy,
+ )
+ return model
diff --git a/src/scmidas/data/quickstart_pbmc_mosaic.h5mu b/src/scmidas/data/quickstart_pbmc_mosaic.h5mu
new file mode 100644
index 0000000..5cc51a8
Binary files /dev/null and b/src/scmidas/data/quickstart_pbmc_mosaic.h5mu differ
diff --git a/src/scmidas/datasets.py b/src/scmidas/datasets.py
new file mode 100644
index 0000000..fecfac8
--- /dev/null
+++ b/src/scmidas/datasets.py
@@ -0,0 +1,45 @@
+"""Bundled example datasets shipped inside the scmidas wheel.
+
+These are toy-sized subsets of real datasets, designed to make the README
+quickstart runnable in under a minute on a single GPU. They are NOT meant
+for benchmarking — see the basics tutorials for full-size data.
+"""
+from __future__ import annotations
+
+from importlib import resources
+from pathlib import Path
+
+import mudata as mu
+
+
+def quickstart_path() -> Path:
+ """Return the on-disk path of the bundled quickstart .h5mu file.
+
+ Returns:
+ Path:
+ Absolute path to ``quickstart_pbmc_mosaic.h5mu`` inside the
+ installed scmidas package.
+ """
+ with resources.as_file(
+ resources.files('scmidas').joinpath('data/quickstart_pbmc_mosaic.h5mu')
+ ) as p:
+ return Path(p)
+
+
+def quickstart() -> mu.MuData:
+ """Load the bundled quickstart MuData (PBMC RNA+ADT mosaic, 1600 cells).
+
+ The dataset is a hand-tuned subset of the WNN PBMC mosaic dataset:
+ 4 batches × 400 cells each (RNA-only, ADT-only, two paired) with
+ 500 RNA HVGs + 224 ADT features, sized so that
+ ``scmidas.integrate(...)`` finishes in roughly one minute on a
+ single mid-range GPU. **It is intended for the quickstart only**;
+ its size and feature count are not appropriate for serious analysis.
+
+ Returns:
+ MuData:
+ A MuData with two modalities (``'rna'``, ``'adt'``) and the
+ following ``obs`` columns at top level:
+ ``'batch'`` and ``'celltype'``.
+ """
+ return mu.read_h5mu(str(quickstart_path()))