Skills › Scientific Agent Skills › scanpy › annotate.py
annotate.py
86 lines · skills/scanpy/scripts/annotate.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Annotate clusters with cell-type labels from a mapping file.
5Maps a cluster column (e.g. leiden) to cell-type names using a JSON or CSV
6mapping, writes the labels into a new obs column, and saves a UMAP and dotplot.
8The mapping file is one of:
9 * JSON : {"0": "CD4 T cells", "1": "B cells", ...}
10 * CSV : two columns ``cluster,cell_type``
12Optionally provide a marker-gene JSON to draw a reference dotplot that helps
13decide the mapping:
14 {"T cells": ["CD3D","CD3E"], "B cells": ["MS4A1","CD79A"]}
16Examples:
17 python annotate.py clustered.h5ad -o annotated.h5ad --mapping celltypes.json
18 python annotate.py clustered.h5ad -o annotated.h5ad --mapping map.csv --cluster-key leiden
19 python annotate.py clustered.h5ad --markers markers.json --cluster-key leiden # dotplot only
20"""
22import argparse
23import json
24import os
26from _common import add_io_args, configure_scanpy, die, info, load_anndata, save_anndata
29def load_mapping(path):
30 if path.lower().endswith(".json"):
31 with open(path) as fh:
32 return {str(k): v for k, v in json.load(fh).items()}
33 import pandas as pd
34 df = pd.read_csv(path, dtype=str, keep_default_na=False)
35 if df.shape[1] < 2:
36 die("CSV mapping needs at least two columns: cluster,cell_type")
37 if df.iloc[:, 0].duplicated().any():
38 die("CSV mapping contains duplicate cluster labels")
39 return {str(k): v for k, v in zip(df.iloc[:, 0], df.iloc[:, 1])}
42def main():
43 p = argparse.ArgumentParser(description=__doc__,
44 formatter_class=argparse.RawDescriptionHelpFormatter)
45 add_io_args(p, default_output="annotated.h5ad")
46 p.add_argument("--cluster-key", default="leiden", help="obs column with clusters (default leiden)")
47 p.add_argument("--mapping", default=None, help="JSON or CSV cluster->cell_type mapping")
48 p.add_argument("--label-key", default="cell_type", help="New obs column name (default cell_type)")
49 p.add_argument("--markers", default=None,
50 help="JSON of {cell_type: [genes]} to draw a reference dotplot")
51 p.add_argument("--no-plots", action="store_true", help="Skip plots")
52 args = p.parse_args()
54 sc = configure_scanpy(figdir=args.figdir)
55 adata = load_anndata(args.input)
56 if args.cluster_key not in adata.obs.columns:
57 die(f"cluster key '{args.cluster_key}' not in obs: {list(adata.obs.columns)}")
59 if args.markers:
60 with open(args.markers) as fh:
61 marker_dict = json.load(fh)
62 present = {ct: [g for g in genes if g in (adata.raw.var_names if adata.raw is not None else adata.var_names)]
63 for ct, genes in marker_dict.items()}
64 present = {ct: g for ct, g in present.items() if g}
65 if present and not args.no_plots:
66 sc.pl.dotplot(adata, present, groupby=args.cluster_key,
67 use_raw=adata.raw is not None, show=False, save="_marker_reference.png")
68 info("Wrote marker reference dotplot")
70 if args.mapping:
71 mapping = load_mapping(args.mapping)
72 adata.obs[args.label_key] = (
73 adata.obs[args.cluster_key].astype(str).map(mapping).fillna("Unknown").astype("category")
74 )
75 info(f"Annotated '{args.label_key}': "
76 + ", ".join(f"{k}={v}" for k, v in adata.obs[args.label_key].value_counts().items()))
77 if not args.no_plots:
78 sc.pl.umap(adata, color=args.label_key, legend_loc="on data",
79 show=False, save="_celltypes.png")
80 save_anndata(adata, args.output)
81 elif not args.markers:
82 die("provide --mapping (to annotate) and/or --markers (for a reference dotplot)")
85if __name__ == "__main__":
86 main()