Skills › Scientific Agent Skills › scanpy › batch_correct.py
batch_correct.py
73 lines · skills/scanpy/scripts/batch_correct.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Batch correction / integration across samples.
5Supports three methods:
6 * harmony : corrects the PCA embedding -> writes obsm['X_pca_harmony'].
7 Embedding integration. Needs harmonypy (uv pip install harmonypy==2.0.2).
8 Follow with: reduce_dimensions.py --use-rep X_pca_harmony
9 * bbknn : batch-balanced kNN graph (replaces sc.pp.neighbors). Then cluster directly.
10 Needs bbknn (uv pip install bbknn).
11 * combat : corrects the expression matrix in place (sc.pp.combat). Built into scanpy.
13Run on a normalized object that already has PCA (harmony/bbknn) computed.
15Examples:
16 python batch_correct.py reduced.h5ad -o integrated.h5ad --method harmony --batch-key sample
17 python batch_correct.py reduced.h5ad -o integrated.h5ad --method bbknn --batch-key sample
18 python batch_correct.py normalized.h5ad -o integrated.h5ad --method combat --batch-key batch
19"""
21import argparse
23from _common import add_io_args, configure_scanpy, die, info, load_anndata, save_anndata, compute_pca, clear_graph, ensure_categories, integrate_harmony
26def main():
27 p = argparse.ArgumentParser(description=__doc__,
28 formatter_class=argparse.RawDescriptionHelpFormatter)
29 add_io_args(p, default_output="integrated.h5ad")
30 p.add_argument("--method", default="harmony", choices=["harmony", "bbknn", "combat"])
31 p.add_argument("--batch-key", required=True, help="obs column identifying batches")
32 args = p.parse_args()
34 sc = configure_scanpy(figdir=args.figdir)
35 adata = load_anndata(args.input)
36 if args.batch_key not in adata.obs.columns:
37 die(f"batch key '{args.batch_key}' not in obs: {list(adata.obs.columns)}")
39 ensure_categories(adata, args.batch_key)
40 if adata.obs[args.batch_key].nunique() < 2:
41 die("batch correction requires at least two batches")
43 if args.method == "harmony":
44 if "X_pca" not in adata.obsm:
45 compute_pca(sc, adata)
46 try:
47 integrate_harmony(adata, args.batch_key)
48 except ImportError:
49 die("harmonypy not installed. Install with: uv pip install harmonypy==2.0.2")
50 clear_graph(adata)
51 info("Wrote obsm['X_pca_harmony']. Next: "
52 "reduce_dimensions.py --use-rep X_pca_harmony")
53 elif args.method == "bbknn":
54 if "X_pca" not in adata.obsm:
55 compute_pca(sc, adata)
56 try:
57 sc.external.pp.bbknn(adata, batch_key=args.batch_key,
58 n_pcs=adata.obsm["X_pca"].shape[1],
59 approx=False, use_faiss=False)
60 except ImportError:
61 die("bbknn not installed. Install with: uv pip install bbknn")
62 sc.tl.umap(adata)
63 info("Built batch-balanced graph + UMAP. Next: cluster.py")
64 elif args.method == "combat":
65 sc.pp.combat(adata, key=args.batch_key)
66 clear_graph(adata, clear_pca=True)
67 info("Corrected expression matrix with ComBat. Re-run reduce_dimensions.py.")
69 save_anndata(adata, args.output)
72if __name__ == "__main__":
73 main()