Add tutorials, benchmark report, and docs

docs/index.md

@@ -6,6 +6,7 @@
 :hidden: true
 :maxdepth: 1
 
+tutorials.md
 api.md
 changelog.md
 contributing.md

docs/notebooks/lung_example.ipynb

@@ -0,0 +1,159 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "wei6g2uemfc",
+   "source": "# Benchmarking lung integration\n\nHere we walkthrough applying integration benchmarking metrics on the lung atlas example from the [scIB paper](https://www.nature.com/articles/s41592-021-01336-8). This mirrors the [scib-metrics tutorial](https://scib-metrics.readthedocs.io/en/stable/notebooks/lung_example.html), but uses scib-rapids for GPU-accelerated computation.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "k1ufrfd0xe",
+   "source": "import time\n\nimport numpy as np\nimport pandas as pd\nimport scanpy as sc\n\nimport scib_rapids\nfrom scib_rapids.nearest_neighbors import NeighborsResults, pynndescent",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "u7l9ngw927b",
+   "source": "## Load and preprocess data",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "y9bsu2y7k9k",
+   "source": "adata = sc.read(\n    \"data/lung_atlas.h5ad\",\n    backup_url=\"https://figshare.com/ndownloader/files/24539942\",\n)\nadata",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "6wp4m7qcn7t",
+   "source": "sc.pp.highly_variable_genes(adata, n_top_genes=2000, flavor=\"cell_ranger\", batch_key=\"batch\")\nsc.tl.pca(adata, n_comps=30, use_highly_variable=True)",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "q7ppcgda4vk",
+   "source": "We subset to the highly variable genes so that each method has the same input.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "36kjgh5fyb4",
+   "source": "adata = adata[:, adata.var.highly_variable].copy()\nadata.obsm[\"Unintegrated\"] = adata.obsm[\"X_pca\"]",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eah0tzhtdzg",
+   "source": "## Run integration methods\n\nHere we run a few embedding-based methods. By focusing on embedding-based methods, we can substantially reduce the runtime of the benchmarking metrics.\n\n### Harmony",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "sfvgtfdgrnk",
+   "source": "from harmony import harmonize\n\nadata.obsm[\"Harmony\"] = harmonize(adata.obsm[\"X_pca\"], adata.obs, batch_key=\"batch\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42bnd5dmqke",
+   "source": "### scVI",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "stz57vrnvmd",
+   "source": "import scvi\n\nscvi.model.SCVI.setup_anndata(adata, layer=\"counts\", batch_key=\"batch\")\nvae = scvi.model.SCVI(adata, gene_likelihood=\"nb\", n_layers=2, n_latent=30)\nvae.train()\nadata.obsm[\"scVI\"] = vae.get_latent_representation()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7ko3cf0gmx4",
+   "source": "### scANVI",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "te2a3x7vzkr",
+   "source": "lvae = scvi.model.SCANVI.from_scvi_model(\n    vae,\n    adata=adata,\n    labels_key=\"cell_type\",\n    unlabeled_category=\"Unknown\",\n)\nlvae.train(max_epochs=20, n_samples_per_label=100)\nadata.obsm[\"scANVI\"] = lvae.get_latent_representation()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0dyu1vtgmkns",
+   "source": "## Compute all metrics\n\nWe define a helper function that computes all metrics for a given embedding and returns a results dictionary.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "wlmt2ehwmo",
+   "source": "def compute_all_metrics(\n    adata,\n    embedding_key,\n    batch_key=\"batch\",\n    label_key=\"cell_type\",\n    pre_integrated_key=\"X_pca\",\n    n_neighbors=90,\n):\n    \"\"\"Compute all scib-rapids metrics for a single embedding.\"\"\"\n    X = np.asarray(adata.obsm[embedding_key], dtype=np.float32)\n    labels = np.asarray(pd.Categorical(adata.obs[label_key]).codes)\n    batch = np.asarray(pd.Categorical(adata.obs[batch_key]).codes)\n\n    # Compute nearest neighbors\n    nn = pynndescent(X, n_neighbors=n_neighbors)\n\n    results = {}\n\n    # Bio conservation\n    results[\"Isolated labels\"] = scib_rapids.isolated_labels(X, labels, batch)\n    kmeans = scib_rapids.nmi_ari_cluster_labels_kmeans(X, labels)\n    results[\"KMeans NMI\"] = kmeans[\"nmi\"]\n    results[\"KMeans ARI\"] = kmeans[\"ari\"]\n    results[\"Silhouette label\"] = scib_rapids.silhouette_label(X, labels)\n    results[\"cLISI\"] = scib_rapids.clisi_knn(nn, labels)\n\n    # Batch correction\n    results[\"Silhouette batch\"] = scib_rapids.silhouette_batch(X, labels, batch)\n    results[\"iLISI\"] = scib_rapids.ilisi_knn(nn, batch)\n    results[\"KBET\"] = scib_rapids.kbet_per_label(nn, batch, labels)\n    results[\"Graph connectivity\"] = scib_rapids.graph_connectivity(nn, labels)\n\n    # PCR comparison (requires pre-integrated embedding)\n    if pre_integrated_key is not None:\n        X_pre = np.asarray(adata.obsm[pre_integrated_key], dtype=np.float32)\n        results[\"PCR comparison\"] = scib_rapids.pcr_comparison(X_pre, X, batch, categorical=True)\n    else:\n        results[\"PCR comparison\"] = 0.0\n\n    return results",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cprkjbq718h",
+   "source": "## Perform the benchmark\n\nRun all metrics on each embedding.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "l56sgynqc7",
+   "source": "embedding_keys = [\"Unintegrated\", \"Harmony\", \"scVI\", \"scANVI\"]\n\nall_results = {}\nfor key in embedding_keys:\n    print(f\"Computing metrics for {key}...\")\n    t0 = time.time()\n    all_results[key] = compute_all_metrics(adata, key)\n    elapsed = time.time() - t0\n    print(f\"  Done in {elapsed:.1f}s\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4imqdpx10pt",
+   "source": "## Visualize the results",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "nbx89tlqvq9",
+   "source": "df = pd.DataFrame(all_results).T\ndf.index.name = \"Embedding\"\n\n# Add metric type annotations\nbio_metrics = [\"Isolated labels\", \"KMeans NMI\", \"KMeans ARI\", \"Silhouette label\", \"cLISI\"]\nbatch_metrics = [\"Silhouette batch\", \"iLISI\", \"KBET\", \"Graph connectivity\", \"PCR comparison\"]\n\n# Compute aggregate scores\ndf[\"Bio conservation\"] = df[bio_metrics].mean(axis=1)\ndf[\"Batch correction\"] = df[batch_metrics].mean(axis=1)\ndf[\"Total\"] = 0.6 * df[\"Bio conservation\"] + 0.4 * df[\"Batch correction\"]\n\ndf.round(4)",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "adwoc8otaih",
+   "source": "df[[\"Bio conservation\", \"Batch correction\", \"Total\"]].sort_values(\"Total\", ascending=False)",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/notebooks/quickstart.ipynb

@@ -0,0 +1,201 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "hfcdfmavktq",
+   "source": "# Quickstart: scib-rapids metrics\n\nThis tutorial demonstrates how to use scib-rapids to compute single-cell integration benchmarking metrics. scib-rapids provides the same metrics as [scib-metrics](https://github.com/YosefLab/scib-metrics), but uses CuPy/RAPIDS for GPU acceleration instead of JAX.\n\nThe API is identical — if you're familiar with scib-metrics, you can use scib-rapids as a drop-in replacement.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "b3qo8n0m9ym",
+   "source": "import numpy as np\nimport scib_rapids\nfrom scib_rapids.nearest_neighbors import NeighborsResults, pynndescent",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0rcwtblzizq",
+   "source": "## Generate example data\n\nWe create a simple toy dataset with 1000 cells, 50 features, 5 cell types, and 3 batches. In practice, `X` would be a PCA embedding from an AnnData object.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "9zqcbsj3ins",
+   "source": "rng = np.random.default_rng(42)\nn_cells, n_features, n_labels, n_batches = 1000, 50, 5, 3\n\nX = rng.normal(size=(n_cells, n_features)).astype(np.float32)\nlabels = rng.integers(0, n_labels, size=n_cells)\nbatch = rng.integers(0, n_batches, size=n_cells)\n\nprint(f\"X: {X.shape}, labels: {np.unique(labels)}, batches: {np.unique(batch)}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ysiihlzae5i",
+   "source": "## Compute nearest neighbors\n\nSeveral metrics (LISI, kBET, graph connectivity, Leiden clustering) require a precomputed k-nearest neighbor graph. scib-rapids provides a `pynndescent` wrapper that returns a `NeighborsResults` dataclass.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "rwf2yqookf",
+   "source": "nn = pynndescent(X, n_neighbors=30)\nprint(f\"NeighborsResults: indices {nn.indices.shape}, distances {nn.distances.shape}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "wxytsr0hga",
+   "source": "## Bio-conservation metrics\n\nThese metrics measure how well biological signal (cell type identity) is preserved after integration.\n\n### Silhouette scores",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "0f3wpmxrfpfe",
+   "source": "# Average Silhouette Width — measures cell type separation\n# Returns a score in [0, 1] (rescaled) where higher = better separation\nasw = scib_rapids.silhouette_label(X, labels)\nprint(f\"Silhouette label (ASW): {asw:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "db5lpk6t86l",
+   "source": "# Isolated labels — score for cell types present in few batches\niso = scib_rapids.isolated_labels(X, labels, batch)\nprint(f\"Isolated labels: {iso:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bppaoyqjq2p",
+   "source": "### Clustering-based metrics (NMI / ARI)",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "2daq4xiyto6",
+   "source": "# K-means clustering agreement with true labels\nkmeans_result = scib_rapids.nmi_ari_cluster_labels_kmeans(X, labels)\nprint(f\"KMeans NMI: {kmeans_result['nmi']:.4f}\")\nprint(f\"KMeans ARI: {kmeans_result['ari']:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "id": "gk87gnmbmgj",
+   "source": "# Leiden clustering agreement (uses the kNN graph)\nleiden_result = scib_rapids.nmi_ari_cluster_labels_leiden(nn, labels)\nprint(f\"Leiden NMI: {leiden_result['nmi']:.4f}\")\nprint(f\"Leiden ARI: {leiden_result['ari']:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3jzyghql9hp",
+   "source": "### cLISI — cell-type Local Inverse Simpson Index",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "kgmn3rzv6un",
+   "source": "# cLISI — how well cell types are preserved in neighborhoods\n# Higher = better cell type preservation (scaled to [0, 1])\nclisi = scib_rapids.clisi_knn(nn, labels)\nprint(f\"cLISI: {clisi:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7lhgb3o3ylo",
+   "source": "## Batch-correction metrics\n\nThese metrics measure how well batch effects have been removed.\n\n### Silhouette batch & BRAS",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "5c1ahtgwqzv",
+   "source": "# Silhouette batch — measures batch mixing within cell types\nsil_batch = scib_rapids.silhouette_batch(X, labels, batch)\nprint(f\"Silhouette batch: {sil_batch:.4f}\")\n\n# BRAS — Batch Removal Adapted Silhouette (uses cosine distance, mean_other)\nbras_score = scib_rapids.bras(X, labels, batch)\nprint(f\"BRAS: {bras_score:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "997uerovcmo",
+   "source": "### iLISI — integration LISI",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "g7mq1ji756g",
+   "source": "# iLISI — how well batches are mixed in neighborhoods\n# Higher = better batch integration (scaled to [0, 1])\nilisi = scib_rapids.ilisi_knn(nn, batch)\nprint(f\"iLISI: {ilisi:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "qs7qfkjvhri",
+   "source": "### kBET",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "7ia6pfojmbi",
+   "source": "# kBET — batch effect test via chi-square statistics\n# acceptance_rate close to 1 means batches are well mixed\nacceptance_rate, test_stats, p_values = scib_rapids.kbet(nn, batch)\nprint(f\"kBET acceptance rate: {acceptance_rate:.4f}\")\n\n# Per-label kBET (uses diffusion distances within each cell type)\nkbet_score = scib_rapids.kbet_per_label(nn, batch, labels)\nprint(f\"kBET per label: {kbet_score:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "uc1tbd9hl3i",
+   "source": "### Graph connectivity",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "q6fqb045uer",
+   "source": "# Graph connectivity — fraction of cells in the largest connected component\n# per cell type subgraph. Higher = better.\ngc = scib_rapids.graph_connectivity(nn, labels)\nprint(f\"Graph connectivity: {gc:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "hmm14u2au94",
+   "source": "### PCR comparison",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "iz8vintxvyn",
+   "source": "# PCR comparison — reduction in batch variance explained by PCA\n# Compares pre- vs post-integration embeddings\nX_post = rng.normal(size=X.shape).astype(np.float32)\npcr = scib_rapids.pcr_comparison(X, X_post, batch, categorical=True)\nprint(f\"PCR comparison: {pcr:.4f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "vro1fk08j3j",
+   "source": "## Summary table\n\nCollect all metrics into a single table.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "a67mgvo76v",
+   "source": "import pandas as pd\n\nresults = {\n    \"Metric\": [\n        \"Silhouette label\", \"Isolated labels\",\n        \"KMeans NMI\", \"KMeans ARI\",\n        \"Leiden NMI\", \"Leiden ARI\",\n        \"cLISI\",\n        \"Silhouette batch\", \"BRAS\",\n        \"iLISI\", \"kBET\", \"kBET per label\",\n        \"Graph connectivity\", \"PCR comparison\",\n    ],\n    \"Type\": [\n        \"Bio conservation\", \"Bio conservation\",\n        \"Bio conservation\", \"Bio conservation\",\n        \"Bio conservation\", \"Bio conservation\",\n        \"Bio conservation\",\n        \"Batch correction\", \"Batch correction\",\n        \"Batch correction\", \"Batch correction\", \"Batch correction\",\n        \"Batch correction\", \"Batch correction\",\n    ],\n    \"Score\": [\n        asw, iso,\n        kmeans_result[\"nmi\"], kmeans_result[\"ari\"],\n        leiden_result[\"nmi\"], leiden_result[\"ari\"],\n        clisi,\n        sil_batch, bras_score,\n        ilisi, acceptance_rate, kbet_score,\n        gc, pcr,\n    ],\n}\n\ndf = pd.DataFrame(results).set_index(\"Metric\")\ndf[\"Score\"] = df[\"Score\"].map(\"{:.4f}\".format)\ndf",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/tutorials.md

@@ -0,0 +1,9 @@
+# Tutorials
+
+```{toctree}
+:hidden: true
+:maxdepth: 1
+
+notebooks/quickstart
+notebooks/lung_example
+```