Skills › Scientific Agent Skills › scanpy › subset.py
subset.py
65 lines · skills/scanpy/scripts/subset.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Subset an AnnData object by cell metadata or genes.
5Filter cells by obs-column values (keep or drop), or restrict to a gene list.
6Useful for isolating a cell type / condition for focused re-analysis.
8Examples:
9 python subset.py annotated.h5ad -o tcells.h5ad --obs cell_type --keep "CD4 T cells" "CD8 T cells"
10 python subset.py annotated.h5ad -o no_doublets.h5ad --obs predicted_doublet --drop True
11 python subset.py data.h5ad -o panel.h5ad --genes CD3D CD14 MS4A1 NKG7
12"""
14import argparse
16from _common import add_io_args, configure_scanpy, die, info, load_anndata, save_anndata, clear_graph
19def main():
20 p = argparse.ArgumentParser(description=__doc__,
21 formatter_class=argparse.RawDescriptionHelpFormatter)
22 add_io_args(p, default_output="subset.h5ad")
23 p.add_argument("--obs", default=None, help="obs column to filter on")
24 p.add_argument("--keep", nargs="+", default=None, help="Values of --obs to keep")
25 p.add_argument("--drop", nargs="+", default=None, help="Values of --obs to drop")
26 p.add_argument("--genes", nargs="+", default=None, help="Restrict to these genes")
27 p.add_argument("--recompute-hvg", action="store_true",
28 help="After subsetting, drop stale embeddings so the next steps recompute cleanly")
29 args = p.parse_args()
31 configure_scanpy(figdir=args.figdir)
32 adata = load_anndata(args.input)
33 n0 = adata.n_obs
35 if args.obs:
36 if args.obs not in adata.obs.columns:
37 die(f"obs column '{args.obs}' not found: {list(adata.obs.columns)}")
38 vals = adata.obs[args.obs].astype(str)
39 if args.keep:
40 adata = adata[vals.isin([str(v) for v in args.keep]), :].copy()
41 elif args.drop:
42 adata = adata[~vals.isin([str(v) for v in args.drop]), :].copy()
43 else:
44 die("provide --keep or --drop with --obs")
45 info(f"Cells {n0} -> {adata.n_obs}")
47 if args.genes:
48 present = [g for g in args.genes if g in adata.var_names]
49 if not present:
50 die("none of the requested genes are present")
51 adata = adata[:, present].copy()
52 info(f"Restricted to {adata.n_vars} genes")
54 if args.recompute_hvg:
55 clear_graph(adata, clear_pca=True)
56 for key in list(adata.var):
57 if key.startswith("highly_variable"):
58 del adata.var[key]
59 info("Cleared standard HVG/embedding/graph results. Recompute HVGs on the existing log-normalized X; do not normalize it again.")
61 save_anndata(adata, args.output)
64if __name__ == "__main__":
65 main()