Add toy scRNA-seq data generator

Author: challagandla

scripts/create_toy_data.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Create a tiny synthetic scRNA-seq AnnData file for smoke tests and tutorials.
+
+The dataset is intentionally simple: three PBMC-like populations, sparse integer
+counts, standard QC gene prefixes, and basic sample metadata. It is not intended
+for biological benchmarking.
+"""
+
+import argparse
+import os
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+from scipy import sparse
+
+
+def make_gene_names(n_background: int = 560) -> list[str]:
+    marker_genes = [
+        "CD3D", "CD3E", "CD3G", "CD4", "IL7R", "CCR7", "SELL",
+        "CD8A", "CD8B", "GZMB", "GZMK", "NKG7", "GNLY", "MS4A1",
+        "CD79A", "CD79B", "CD14", "LYZ", "S100A8", "S100A9",
+        "FCGR3A", "MS4A7", "FCER1A", "CST3", "LILRA4", "IL3RA",
+        "PPBP", "PF4", "HBB", "HBA1", "HBA2",
+    ]
+    mt_genes = [f"MT-{name}" for name in ["ND1", "ND2", "CO1", "CO2", "ATP6", "CYB"]]
+    ribo_genes = [f"RPL{i}" for i in range(3, 13)] + [f"RPS{i}" for i in range(3, 13)]
+    background = [f"GENE{i:04d}" for i in range(n_background)]
+    return marker_genes + mt_genes + ribo_genes + background
+
+
+def create_toy_adata(n_cells: int, seed: int) -> ad.AnnData:
+    rng = np.random.default_rng(seed)
+    genes = make_gene_names()
+    n_genes = len(genes)
+
+    cell_types = np.array(["CD4 T cell", "B cell", "Classical Mono"])
+    probs = np.array([0.45, 0.25, 0.30])
+    assigned = rng.choice(cell_types, size=n_cells, p=probs)
+
+    base = rng.negative_binomial(n=2, p=0.92, size=(n_cells, n_genes)).astype(np.int32)
+    gene_index = {g: i for i, g in enumerate(genes)}
+
+    signatures = {
+        "CD4 T cell": ["CD3D", "CD3E", "CD3G", "CD4", "IL7R", "CCR7"],
+        "B cell": ["MS4A1", "CD79A", "CD79B"],
+        "Classical Mono": ["CD14", "LYZ", "S100A8", "S100A9"],
+    }
+    for ct, markers in signatures.items():
+        rows = np.where(assigned == ct)[0]
+        cols = [gene_index[g] for g in markers]
+        base[np.ix_(rows, cols)] += rng.poisson(lam=4, size=(len(rows), len(cols))).astype(np.int32)
+
+    # Add a few mitochondrial and ribosomal counts so QC plots are meaningful.
+    mt_cols = [i for i, g in enumerate(genes) if g.startswith("MT-")]
+    ribo_cols = [i for i, g in enumerate(genes) if g.startswith(("RPL", "RPS"))]
+    base[:, mt_cols] += rng.poisson(lam=1, size=(n_cells, len(mt_cols))).astype(np.int32)
+    base[:, ribo_cols] += rng.poisson(lam=1, size=(n_cells, len(ribo_cols))).astype(np.int32)
+
+    obs = pd.DataFrame(
+        {
+            "sample": ["toy_pbmc"] * n_cells,
+            "batch": ["batch1"] * n_cells,
+            "condition": ["control"] * n_cells,
+            "cell_type_truth": assigned,
+        },
+        index=[f"toy_cell_{i:04d}" for i in range(n_cells)],
+    )
+    var = pd.DataFrame(index=genes)
+    adata = ad.AnnData(X=sparse.csr_matrix(base), obs=obs, var=var)
+    adata.var_names_make_unique()
+    return adata
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Create toy scRNA-seq .h5ad data")
+    parser.add_argument("--out", default="tests/data/toy_pbmc_raw.h5ad")
+    parser.add_argument("--n-cells", type=int, default=300)
+    parser.add_argument("--seed", type=int, default=7)
+    args = parser.parse_args()
+
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    adata = create_toy_adata(n_cells=args.n_cells, seed=args.seed)
+    adata.write_h5ad(args.out)
+    print(f"Wrote {args.out}: {adata.n_obs} cells x {adata.n_vars} genes")
+
+
+if __name__ == "__main__":
+    main()