# /// script # requires-python = ">=3.10" # dependencies = ["anndata>=0.10", "h5py", "numpy", "pandas", "scipy"] # /// r"""Prepare a Helical-ready, per-sample subsample of the Chu et al. 2024 CRC atlas. The atlas (671,192 cells, ~200 donors, 9+ source studies) is not in CELLxGENE format, so the layout is discovered rather than assumed. Run it in two steps: # 1. See what's in the file: obs columns, which are sample-level, where counts live uv run prepare_crc_atlas.py --inspect # 2. Subsample Pelka et al. tumour samples, 200 cells each, recovering counts from log-normalised X uv run prepare_crc_atlas.py --counts lognorm:X \ --filter Dataset=Pelka.K Class=T --sample-col Sample --patient-col Patient \ --cells-per-sample 200 --min-cells 100 \ --patient-labels MMRstatus Site Grade \ --sample-field 'sort_code=^[^_]+_T_\d+_\d+_(\d+)_' \ --summarise MMRstatus_patient sort_code Site_patient \ --out data/pelka_crc_tumour_200.h5ad Only the sampled rows are read from disk, so memory stays low. """ import argparse import itertools import json import re import sys from pathlib import Path import anndata as ad import h5py import numpy as np import pandas as pd import scipy.sparse as sp try: # anndata >= 0.11 from anndata.io import read_elem, sparse_dataset except ImportError: # anndata 0.10 from anndata.experimental import read_elem, sparse_dataset DEFAULT_INPUT = Path.home() / "data/chu2024_crc/adata_AllAnnotated.h5ad" SOURCE = "https://doi.org/10.6084/m9.figshare.25323397" DOI = "https://doi.org/10.1038/s43018-024-00807-z" SAMPLE_COL_GUESSES = ["Sample", "sample", "sample_id", "SampleID", "orig.ident"] ENSEMBL = re.compile(r"^ENS[A-Z]*G\d") COUNT_LAYERS = ["counts", "count", "raw_counts"] # conventional names for a raw-counts layer, checked first MISSING = {"", "nan", "None", "NA"} # missing values after astype(str) MAX_LEVELS = 200 # obs columns with more distinct values than this are dropped (barcodes, QC floats) # --- reading ------------------------------------------------------------------------- def read_frame(f: h5py.File, key: str) -> pd.DataFrame: try: return read_elem(f[key]) except Exception: # files written by old anndata versions a = ad.read_h5ad(f.filename, backed="r") return (a.obs if key == "obs" else a.var).copy() def matrix_candidates(f: h5py.File) -> list[str]: layers = list(f.get("layers", {}).keys()) named = [k for k in COUNT_LAYERS if k in layers] keys = [f"layers/{k}" for k in named + [k for k in layers if k not in named]] keys += ["raw/X"] if "raw" in f and "X" in f["raw"] else [] return keys + ["X"] def read_rows(f: h5py.File, key: str, rows: np.ndarray) -> sp.csr_matrix: """Rows (sorted) of a sparse or dense matrix, as CSR.""" node = f[key] if isinstance(node, h5py.Group): return sparse_dataset(node)[rows].tocsr() blocks = [sp.csr_matrix(node[rows[i:i + 2000]]) for i in range(0, len(rows), 2000)] return sp.vstack(blocks).tocsr() def counts_report(f: h5py.File, key: str, n: int) -> dict: shape = f[key].attrs.get("shape", getattr(f[key], "shape", None)) n_obs = int(shape[0]) sample = np.sort(np.random.default_rng(0).choice(n_obs, min(n, n_obs), replace=False)) X = read_rows(f, key, sample) d = X.data return { "shape": [int(s) for s in shape], "storage": "sparse" if isinstance(f[key], h5py.Group) else "dense", "frac_nonint": float(np.mean(np.abs(d - np.round(d)) > 1e-6)) if d.size else 0.0, "max": float(d.max()) if d.size else 0.0, "min": float(d.min()) if d.size else 0.0, "median_total_per_cell": float(np.median(np.asarray(X.sum(axis=1)).ravel())), } def looks_like_counts(report: dict) -> bool: return report["frac_nonint"] < 0.01 and report["min"] >= 0 and report["max"] > 50 def recover_counts(X: sp.csr_matrix, n_count: np.ndarray, target_sum: float): """Undo Seurat LogNormalize (log1p(counts / nCount * target_sum)), row by row. Uses nCount when that gives whole numbers. Otherwise (e.g. Pelka.K, where nCount_RNA does not match the matrix) divides each cell by its smallest non-zero value, which is exact whenever the cell has at least one gene with a count of 1 -- true for UMI data. """ nnz = np.diff(X.indptr) row_of = np.repeat(np.arange(X.shape[0]), nnz) e = np.expm1(X.data) a = e * (n_count / target_sum)[row_of] bad = np.bincount(row_of, weights=np.abs(a - np.round(a)) > 0.01, minlength=X.shape[0]) > 0 mins = np.ones(X.shape[0]) mins[nnz > 0] = np.minimum.reduceat(e, X.indptr[:-1][nnz > 0]) X.data = np.where(bad[row_of], e / mins[row_of], a) # in place: same sparsity, no matrix copies return X, int((~bad).sum()), int(bad.sum()) # --- metadata helpers --------------------------------------------------------------- def sample_level_columns(obs: pd.DataFrame, sample_col: str) -> list[str]: """Columns with exactly one value per sample (patient/sample metadata).""" g = obs.groupby(sample_col, observed=True) return [c for c in obs.columns if c != sample_col and g[c].nunique(dropna=False).max() == 1] def guess_sample_col(obs: pd.DataFrame) -> str | None: return next((c for c in SAMPLE_COL_GUESSES if c in obs), None) def parse_filters(items: list[str]) -> list[tuple[str, list[str]]]: out = [] for item in items: col, _, vals = item.partition("=") if not vals: sys.exit(f"--filter expects COLUMN=VALUE[,VALUE...], got '{item}'") out.append((col, vals.split(","))) return out def inspect(path: Path, n_rows: int) -> None: with h5py.File(path, "r") as f: obs, var = read_frame(f, "obs"), read_frame(f, "var") print(f"{path.name}: {obs.shape[0]:,} cells x {var.shape[0]:,} genes\n") print("== Matrices (pick the one that looks like raw counts: frac_nonint ~0, large totals)") for key in matrix_candidates(f): try: r = counts_report(f, key, n_rows) verdict = "raw counts" if looks_like_counts(r) else "not raw counts" print(f" {key}: {verdict} {r}") except Exception as e: print(f" {key}: unreadable ({e})") names = var.index.astype(str) print(f"\n== Genes: {names.str.match(ENSEMBL).mean():.1%} Ensembl-like, " f"unique={names.is_unique}; first: {list(names[:8])}") print(f" var columns: {list(var.columns)}") print("\n== obs columns (n distinct, top values)") for c in obs.columns: vc = obs[c].astype(str).value_counts() top = ", ".join(f"{k}={v:,}" for k, v in vc.head(6).items()) print(f" {c:28s} {len(vc):>7,} {top}") sample_col = guess_sample_col(obs) if sample_col: cols = sample_level_columns(obs, sample_col) print(f"\n== Sample-level columns (one value per '{sample_col}', {obs[sample_col].nunique()} samples):") print(" " + ", ".join(cols)) else: print(f"\nNo sample column among {SAMPLE_COL_GUESSES}; pass --sample-col") # --- main ----------------------------------------------------------------------------- def main() -> None: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--input", type=Path, default=DEFAULT_INPUT) p.add_argument("--out", type=Path, default=Path("data/chu_crc_subsample.h5ad")) p.add_argument("--inspect", action="store_true", help="describe the file and exit") p.add_argument("--inspect-rows", type=int, default=2000, help="cells sampled per matrix in --inspect") p.add_argument("--counts", help="matrix holding raw counts, e.g. X, raw/X or layers/counts (default: auto). " "'lognorm:X' recovers counts from a log-normalised matrix") p.add_argument("--ncount-col", default="nCount_RNA", help="obs column with library sizes, for lognorm:") p.add_argument("--target-sum", type=float, default=1e4, help="LogNormalize scale factor, for lognorm:") p.add_argument("--patient-labels", nargs="*", default=[], help="per-cell columns to resolve to one value per patient (ignoring NA/Normal)") p.add_argument("--sample-field", nargs="*", default=[], metavar="NAME=REGEX", help="derive a column from the sample name (first regex group)") p.add_argument("--sample-col", help="unit to subsample by (default: guessed)") p.add_argument("--patient-col", help="patient/donor column, kept for patient-level CV splits") p.add_argument("--filter", nargs="*", default=[], metavar="COL=V1,V2", help="keep only matching cells") p.add_argument("--cells-per-sample", type=int, default=200) p.add_argument("--min-cells", type=int, default=100, help="drop samples with fewer cells (after filtering)") p.add_argument("--symbol-col", help="var column with gene symbols (default: var index)") p.add_argument("--summarise", nargs="*", default=[], help="sample-level columns to cross-tabulate") p.add_argument("--seed", type=int, default=0) args = p.parse_args() if not args.input.exists(): sys.exit(f"{args.input} not found; download it from {SOURCE}") if args.inspect: return inspect(args.input, args.inspect_rows) with h5py.File(args.input, "r") as f: obs, var = read_frame(f, "obs"), read_frame(f, "var") sample_col = args.sample_col or guess_sample_col(obs) if sample_col not in obs: sys.exit(f"Sample column '{sample_col}' not found; run --inspect and pass --sample-col") # --- pick the counts matrix --------------------------------------------------- lognorm = bool(args.counts and args.counts.startswith("lognorm:")) if args.counts: counts_key = args.counts.removeprefix("lognorm:") else: reports = {} for k in matrix_candidates(f): # stop at the first match; each check reads 500 random rows reports[k] = counts_report(f, k, 500) is_count_layer = k.removeprefix("layers/") in COUNT_LAYERS if looks_like_counts(reports[k]): if is_count_layer: print(f"Found '{k}' with raw counts") counts_key = k break if is_count_layer: print(f"Warning: '{k}' exists but does not look like raw counts " f"(frac_nonint={reports[k]['frac_nonint']:.3f}, max={reports[k]['max']:.1f}); trying others") else: sys.exit(f"No matrix looks like raw counts: {json.dumps(reports, indent=1)}\nPass --counts") gene_var = read_frame(f, "raw/var") if counts_key.startswith("raw/") else var genes = gene_var.index print(f"Full atlas: {obs.shape[0]:,} cells; counts from '{counts_key}' ({len(genes):,} genes)") # --- filter + subsample per sample -------------------------------------------- mask = np.ones(len(obs), bool) for col, vals in parse_filters(args.filter): if col not in obs: sys.exit(f"Filter column '{col}' not in obs") mask &= obs[col].astype(str).isin(vals).to_numpy() print(f"Filter {col} in {vals}: {mask.sum():,} cells left") sizes = obs.loc[mask, sample_col].value_counts() sizes = sizes[sizes > 0] # categorical value_counts also lists filtered-out samples small = set(sizes[sizes < args.min_cells].index) if small: print(f"Dropping {len(small)} samples with < {args.min_cells} cells") mask &= ~obs[sample_col].isin(small).to_numpy() rng = np.random.default_rng(args.seed) eligible = np.flatnonzero(mask) codes = pd.Categorical(obs[sample_col].to_numpy()[eligible]).codes keep = np.split(eligible[np.argsort(codes, kind="stable")], np.cumsum(np.bincount(codes))[:-1]) # Uniform random cells per sample, so each sample's composition stays representative. rows = np.sort(np.concatenate([r if len(r) <= args.cells_per_sample else rng.choice(r, args.cells_per_sample, replace=False) for r in keep])) print(f"Sampled {len(rows):,} cells from {len(keep)} samples") X = read_rows(f, counts_key, rows) if lognorm: X, n_lib, n_min = recover_counts(X, obs[args.ncount_col].to_numpy(float)[rows], args.target_sum) print(f"Recovered counts from log-normalised '{counts_key}': {n_lib:,} cells via {args.ncount_col}, " f"{n_min:,} via smallest non-zero value") X = X.astype(np.float32, copy=False) # --- counts must be (near-)integers ------------------------------------------------ rounded = np.round(X.data) frac = np.abs(X.data - rounded) n_nonint = int((frac > 1e-6).sum()) if n_nonint: max_off = frac.max() if max_off > 0.5 or (X.data < 0).any(): sys.exit(f"{n_nonint:,} non-integer values in '{counts_key}': looks normalised; pass --counts") print(f"Rounding {n_nonint:,} near-integer values (max offset {max_off:.4f})") X.data = rounded X.eliminate_zeros() # --- genes: symbols only, unique --------------------------------------------------- if args.symbol_col: symbols = gene_var[args.symbol_col].astype(str).to_numpy() else: symbols = genes.astype(str).to_numpy() good = ~pd.Series(symbols).str.match(ENSEMBL).to_numpy() & ~np.isin(symbols, list(MISSING)) totals = np.asarray(X.sum(axis=0)).ravel() order = pd.DataFrame({"symbol": symbols, "total": totals})[good].sort_values("total", ascending=False) dup = order["symbol"].duplicated() cols = np.sort(order.index[~dup].to_numpy()) print(f"Genes: kept {len(cols):,} of {len(symbols):,} " f"({int((~good).sum()):,} Ensembl-only/empty, {int(dup.sum()):,} duplicate symbols)") new_var = pd.DataFrame(index=symbols[cols]) new_var["original_name"] = genes.astype(str).to_numpy()[cols] # --- obs: drop per-cell columns with too many values (barcodes, QC floats) ---------- sub = obs.iloc[rows] keep_cols = [c for c in obs.columns if c in {sample_col, args.patient_col} or sub[c].nunique() <= MAX_LEVELS] new_obs = sub[keep_cols] new_obs.index = new_obs.index.astype(str) # Per-cell labels that are really patient-level (MMRstatus is NA for ~half the cells # of each Pelka sample): take the patient's single non-missing value. if args.patient_labels: if not args.patient_col: sys.exit("--patient-labels needs --patient-col") in_subset = obs[args.patient_col].isin(new_obs[args.patient_col].unique()) patients = obs.loc[in_subset, args.patient_col].astype(str) for c in args.patient_labels: vals = pd.DataFrame({args.patient_col: patients, c: obs.loc[in_subset, c].astype(str)}) vals = vals[~vals[c].isin(MISSING | {"Normal"})] per_patient = vals.groupby(args.patient_col)[c].agg(lambda s: s.iloc[0] if s.nunique() == 1 else "ambiguous") new_obs[f"{c}_patient"] = new_obs[args.patient_col].astype(str).map(per_patient).fillna("NA") print(f"{c}_patient: " + ", ".join(f"{k}={v}" for k, v in new_obs.groupby(args.patient_col, observed=True)[f"{c}_patient"].first().value_counts().items()) + " patients") for item in args.sample_field: name, _, pattern = item.partition("=") new_obs[name] = new_obs[sample_col].astype(str).str.extract(pattern, expand=False).fillna("NA") print(f"{name} from sample names: {new_obs.groupby(sample_col, observed=True)[name].first().value_counts().to_dict()} samples") for c in new_obs.columns: if isinstance(new_obs[c].dtype, pd.CategoricalDtype): new_obs[c] = new_obs[c].cat.remove_unused_categories() elif new_obs[c].dtype == object: new_obs[c] = new_obs[c].astype(str) print(f"Kept obs columns: {keep_cols}") out = ad.AnnData(X=X[:, cols], obs=new_obs, var=new_var) out.uns["provenance"] = json.dumps({ "source": SOURCE, "publication": DOI, "counts_from": args.counts or counts_key, "filters": args.filter, "sample_col": sample_col, "cells_per_sample": args.cells_per_sample, "min_cells": args.min_cells, "seed": args.seed, }) args.out.parent.mkdir(parents=True, exist_ok=True) out.write_h5ad(args.out, compression="gzip") print(f"\nWrote {args.out}: {out.n_obs:,} cells x {out.n_vars:,} genes, {args.out.stat().st_size / 1e6:.1f} MB") # --- per-sample summary -------------------------------------------------------------- meta = sample_level_columns(new_obs, sample_col) summary = new_obs.groupby(sample_col, observed=True)[meta].first() summary["n_cells"] = new_obs[sample_col].value_counts() if args.patient_col and args.patient_col in new_obs: n_per_patient = summary.groupby(args.patient_col).size() print(f"{len(n_per_patient)} patients; {int((n_per_patient > 1).sum())} with more than one sample") path = args.out.with_suffix(".samples.csv") summary.to_csv(path) print(f"Wrote {path} ({len(summary)} samples); sample-level columns: {meta}") xtab_cols = [c for c in args.summarise if c in summary] for a, b in itertools.combinations(xtab_cols, 2): print(f"\nSamples by {a} x {b}:\n" + pd.crosstab(summary[a], summary[b], dropna=False).to_string()) if __name__ == "__main__": main()