Skills › Scientific Agent Skills › scanpy › pseudobulk.py
pseudobulk.py
91 lines · skills/scanpy/scripts/pseudobulk.py at 5f2bc03
1#!/usr/bin/env python3
2"""Sum raw counts by biological sample and cell type for downstream DE.
4Exports genes x pseudobulk samples and aligned metadata. Transpose counts for
5PyDESeq2's samples x genes API. Analyze one cell type at a time with independent
6biological replicates; cells and sample x cell-type rows are not new donors.
8Example:
9 python pseudobulk.py annotated.h5ad --by sample cell_type --metadata condition donor --out-prefix results/pb
10"""
12import argparse
13import os
15from _common import configure_scanpy, die, info, load_anndata, validate_counts
18def aggregate_counts(adata, by, layer="counts", metadata=(), min_cells=1):
19 """Collision-free groups, exact layer selection, constant group metadata."""
20 import numpy as np
21 import pandas as pd
22 import scanpy as sc
24 if min_cells < 1 or not by or len(set(by)) != len(by):
25 die("--by must contain unique grouping columns and --min-cells must be positive")
26 for key in list(by) + list(metadata):
27 if key not in adata.obs or adata.obs[key].isna().any():
28 die(f"obs column '{key}' must exist and contain no missing values")
29 if layer is not None and layer not in adata.layers:
30 die(f"raw-count layer '{layer}' is missing; no fallback to X is permitted")
31 matrix = adata.X if layer is None else adata.layers[layer]
32 validate_counts(matrix, "pseudobulk counts", require_nonzero_cells=False)
33 if not adata.var_names.is_unique:
34 die("gene identifiers must be unique before aggregation")
35 # Scanpy joins multiple keys with underscores; ambiguous labels can collide.
36 # Aggregate by opaque IDs and retain the original grouping columns separately.
37 codes, _ = pd.factorize(pd.MultiIndex.from_frame(adata.obs[list(by)]), sort=False)
38 ids = pd.Index([f"pb{i:06d}" for i in range(int(codes.max()) + 1)], name="sample_id")
39 cols = list(dict.fromkeys(list(by) + list(metadata)))
40 rows = []
41 for code, sample_id in enumerate(ids):
42 group = adata.obs.iloc[np.flatnonzero(codes == code)]
43 if any(group[key].nunique(dropna=False) != 1 for key in metadata):
44 die(f"metadata is not constant within pseudobulk group {sample_id}")
45 row = group.iloc[0][cols].to_dict()
46 row["n_cells"] = len(group)
47 rows.append(row)
48 meta = pd.DataFrame(rows, index=ids)
49 # No copy of expression and no mutation of the source object.
50 import anndata as ad
51 work = ad.AnnData(matrix, obs=pd.DataFrame({"group": pd.Categorical(ids[codes])},
52 index=adata.obs_names), var=adata.var.copy())
53 pb = sc.get.aggregate(work, by="group", func="sum")
54 mat = pb.layers["sum"] # aggregate returns layers, not an expression matrix in X
55 arr = mat.toarray() if hasattr(mat, "toarray") else np.asarray(mat)
56 ordered = pb.obs["group"].astype(str).to_numpy()
57 meta = meta.loc[ordered]
58 keep = meta["n_cells"].to_numpy() >= min_cells
59 if not keep.any():
60 die("no pseudobulk groups meet --min-cells")
61 counts = pd.DataFrame(arr[keep].T, index=adata.var_names,
62 columns=meta.index[keep])
63 counts.index.name = "gene_id"
64 return counts, meta.iloc[np.flatnonzero(keep)]
67def main():
68 p = argparse.ArgumentParser(description=__doc__)
69 p.add_argument("input", help="Input .h5ad with full-gene raw counts")
70 p.add_argument("--by", nargs="+", required=True, help="Sample and cell-type obs columns")
71 source = p.add_mutually_exclusive_group()
72 source.add_argument("--layer", default="counts", help="Raw counts layer (default counts)")
73 source.add_argument("--use-x", action="store_true", help="Explicitly declare X as raw counts")
74 p.add_argument("--func", choices=["sum"], default="sum", help="Count-based DE requires sums")
75 p.add_argument("--metadata", nargs="+", default=[], help="Group-constant donor/condition/design columns")
76 p.add_argument("--min-cells", type=int, default=1, help="Minimum contributing cells per profile")
77 p.add_argument("--out-prefix", default="pseudobulk")
78 args = p.parse_args()
79 configure_scanpy()
80 adata = load_anndata(args.input)
81 counts, meta = aggregate_counts(adata, args.by, None if args.use_x else args.layer,
82 args.metadata, args.min_cells)
83 os.makedirs(os.path.dirname(os.path.abspath(args.out_prefix)), exist_ok=True)
84 counts.to_csv(f"{args.out_prefix}_counts.csv")
85 meta.to_csv(f"{args.out_prefix}_samples.csv")
86 info(f"Wrote {counts.shape[0]} genes x {counts.shape[1]} pseudobulk profiles")
87 info("Transpose counts for PyDESeq2; preserve IDs as strings, align metadata, and fit each cell type with biological replicates.")
90if __name__ == "__main__":
91 main()