Skills › Scientific Agent Skills › bulk-rnaseq › build_counts_matrix.py
build_counts_matrix.py
310 lines · skills/bulk-rnaseq/scripts/build_counts_matrix.py at 5f2bc03
1#!/usr/bin/env python3
2"""Assemble a gene-level counts matrix from RNA-seq quantification output.
4Bridges the upstream (Salmon / STAR / featureCounts) and downstream (PyDESeq2)
5halves of a bulk RNA-seq pipeline. Writes:
7 counts.csv genes x samples, INTEGER counts (never TPM/FPKM)
8 metadata_template.csv one row per sample (index=sample) to fill in for DE
9 counts_provenance.json input hashes, count mode and sample order
11Hand both to the `pydeseq2` skill. counts.csv stays genes x samples; the
12pydeseq2 loader transposes to samples x genes.
14Examples
15--------
16# Salmon: per-sample quant dirs (each containing quant.sf) + a tx2gene map
17python build_counts_matrix.py --from salmon \
18 --quant-dir quant/ --tx2gene tx2gene.tsv --output-dir counts/
20# STAR --quantMode GeneCounts: a dir of *.ReadsPerGene.out.tab files
21python build_counts_matrix.py --from star \
22 --quant-dir star/ --strandedness reverse --output-dir counts/
24# featureCounts: the combined matrix it wrote
25python build_counts_matrix.py --from featurecounts \
26 --counts-file counts/featurecounts.txt --output-dir counts/
28Requires: pandas and NumPy. Salmon mode also needs pytximport.
29"""
30from __future__ import annotations
32import argparse
33import hashlib
34import json
35import sys
36from importlib.metadata import version
37from pathlib import Path
39import numpy as np
40import pandas as pd
42from _tabular import validate_header
44# STAR ReadsPerGene.out.tab column index per strandedness (col 0 is the gene id).
45STAR_STRAND_COL = {"unstranded": 1, "forward": 2, "reverse": 3}
48def _clean_sample_name(name: str) -> str:
49 """Strip common aligner suffixes/extensions from a file or column name."""
50 name = Path(str(name)).name
51 for suf in (
52 ".ReadsPerGene.out.tab",
53 ".Aligned.sortedByCoord.out.bam",
54 ".Aligned.out.bam",
55 ".bam",
56 ".sf",
57 ):
58 if name.endswith(suf):
59 name = name[: -len(suf)]
60 return name.rstrip(".")
63def _validate_counts(counts: pd.DataFrame, *, estimated: bool = False) -> pd.DataFrame:
64 """Reject ambiguous labels and malformed values before any integer conversion."""
65 if counts.empty:
66 raise ValueError("Counts matrix is empty")
67 for label, values in (("gene", counts.index), ("sample", counts.columns)):
68 if values.has_duplicates or values.isna().any() or any(not str(v).strip() for v in values):
69 raise ValueError(f"Empty or duplicate {label} IDs in counts matrix")
70 numeric = counts.apply(pd.to_numeric, errors="raise")
71 values = numeric.to_numpy()
72 if not np.isfinite(values).all() or (values < 0).any():
73 raise ValueError("Counts must be finite and non-negative")
74 if (values >= 2**63).any():
75 raise ValueError("Counts exceed the supported int64 range")
76 if not estimated and (values != np.floor(values)).any():
77 raise ValueError("Fractional counts require an explicit estimated-count workflow; refusing truncation")
78 return numeric.round().astype("int64")
81def _file_record(path: Path) -> dict:
82 digest = hashlib.sha256()
83 with path.open("rb") as handle:
84 for chunk in iter(lambda: handle.read(1024 * 1024), b""):
85 digest.update(chunk)
86 return {"path": str(path), "sha256": digest.hexdigest()}
89def build_from_salmon(quant_dir: Path, tx2gene: Path,
90 ignore_transcript_version: bool = False) -> pd.DataFrame:
91 """Aggregate per-sample Salmon quant.sf to gene level via pytximport.
93 Uses counts_from_abundance='length_scaled_tpm' for full-length bulk
94 gene-level DE and rounds the resulting estimated counts to integers,
95 which PyDESeq2 requires. See references/counts-and-handoff.md.
96 """
97 try:
98 from pytximport import tximport
99 except ImportError:
100 sys.exit("Salmon mode needs pytximport. Install with: uv pip install pytximport")
102 if tx2gene is None:
103 sys.exit("--tx2gene is required for --from salmon (columns: transcript_id, gene_id)")
105 # Discover per-sample quant.sf files (sample name = parent directory name).
106 sf_files = sorted(quant_dir.glob("*/quant.sf"))
107 if not sf_files:
108 # Fall back to a flat layout: *.sf directly in quant_dir.
109 sf_files = sorted(quant_dir.glob("*.sf"))
110 if not sf_files:
111 sys.exit(f"No quant.sf found under {quant_dir} (expected <sample>/quant.sf)")
113 sample_names = [
114 f.parent.name if f.name == "quant.sf" else _clean_sample_name(f.name)
115 for f in sf_files
116 ]
117 if len(sample_names) != len(set(sample_names)):
118 raise ValueError("Duplicate sample IDs after filename normalization")
119 validate_header(tx2gene, delimiter=None)
120 mapping = pd.read_csv(tx2gene, sep=None, engine="python", dtype=str, keep_default_na=False)
121 if not {"transcript_id", "gene_id"}.issubset(mapping.columns):
122 raise ValueError("tx2gene needs transcript_id and gene_id headers")
123 mapping = mapping[["transcript_id", "gene_id"]].copy()
124 if mapping.empty or mapping.eq("").any().any():
125 raise ValueError("tx2gene has empty transcript/gene IDs")
126 if ignore_transcript_version:
127 mapping["transcript_id"] = mapping["transcript_id"].str.replace(r"\.\d+$", "", regex=True)
128 if mapping["transcript_id"].duplicated().any():
129 raise ValueError("Duplicate transcript IDs in tx2gene (possibly after version stripping)")
130 reference_ids = None
131 for path in sf_files:
132 validate_header(path, delimiter="\t")
133 quant = pd.read_csv(path, sep="\t", dtype={"Name": str}, keep_default_na=False)
134 required = {"Name", "Length", "EffectiveLength", "TPM", "NumReads"}
135 if not required.issubset(quant.columns):
136 raise ValueError(f"Invalid Salmon quant.sf columns: {path}")
137 names = quant["Name"]
138 if ignore_transcript_version:
139 names = names.str.replace(r"\.\d+$", "", regex=True)
140 if names.duplicated().any() or names.eq("").any():
141 raise ValueError(f"Empty or duplicate transcript IDs: {path}")
142 numeric = quant[["Length", "EffectiveLength", "TPM", "NumReads"]].apply(pd.to_numeric)
143 if not np.isfinite(numeric.to_numpy()).all() or (numeric < 0).any().any():
144 raise ValueError(f"Invalid Salmon numeric values: {path}")
145 if numeric["NumReads"].sum() <= 0:
146 raise ValueError(f"No quantified fragments: {path}")
147 if reference_ids is not None and not names.equals(reference_ids):
148 raise ValueError("Salmon transcript IDs/order differ; use one reference and quantifier run convention")
149 reference_ids = names
150 missing = set(names) - set(mapping["transcript_id"])
151 if missing:
152 raise ValueError(f"tx2gene missing {len(missing)} quantified transcript IDs; e.g. {sorted(missing)[:5]}")
153 print(f"Salmon: {len(sf_files)} samples -> {sample_names}")
155 txi = tximport(
156 [str(f) for f in sf_files],
157 data_type="salmon",
158 transcript_gene_map=mapping,
159 counts_from_abundance="length_scaled_tpm",
160 ignore_transcript_version=ignore_transcript_version,
161 ignore_after_bar=False,
162 output_type="anndata",
163 return_data=True,
164 )
165 # AnnData: obs=samples, var=genes, X=samples x genes. Transpose to genes x samples.
166 counts = txi.to_df().T
167 counts.columns = sample_names
168 counts.index.name = "gene_id"
169 counts = _validate_counts(counts, estimated=True)
170 counts.attrs["provenance"] = {
171 "source": "salmon", "counts_from_abundance": "length_scaled_tpm",
172 "assay": "full_length_bulk", "rounding": "nearest_integer_ties_to_even",
173 "ignore_transcript_version": ignore_transcript_version,
174 "pytximport_version": version("pytximport"), "tx2gene": _file_record(tx2gene),
175 "inputs": [_file_record(f) for f in sf_files],
176 }
177 return counts
180def build_from_star(quant_dir: Path, strandedness: str) -> pd.DataFrame:
181 """Combine STAR *.ReadsPerGene.out.tab files into a gene x sample matrix."""
182 col = STAR_STRAND_COL[strandedness]
183 tabs = sorted(quant_dir.glob("*ReadsPerGene.out.tab"))
184 if not tabs:
185 sys.exit(f"No *ReadsPerGene.out.tab found under {quant_dir}")
186 print(f"STAR: {len(tabs)} samples, strandedness={strandedness} (column {col})")
188 series = {}
189 reference_genes = None
190 for tab in tabs:
191 # First 4 rows are summary stats (N_unmapped, N_multimapping, ...).
192 df = pd.read_csv(tab, sep="\t", header=None, dtype={0: str}, keep_default_na=False)
193 if df.shape[1] != 4 or df.iloc[:4, 0].tolist() != ["N_unmapped", "N_multimapping", "N_noFeature", "N_ambiguous"]:
194 raise ValueError(f"Invalid STAR GeneCounts header/layout: {tab}")
195 df = df.iloc[4:]
196 sample = _clean_sample_name(tab.name)
197 if sample in series:
198 raise ValueError(f"Duplicate sample ID after filename normalization: {sample}")
199 s = pd.Series(df[col].values, index=df[0].values)
200 _validate_counts(s.to_frame(sample))
201 if reference_genes is not None and set(s.index) != reference_genes:
202 raise ValueError("STAR gene sets differ; check annotation/reference instead of filling missing genes with zero")
203 reference_genes = set(s.index)
204 series[sample] = s
206 counts = _validate_counts(pd.DataFrame(series))
207 counts.index.name = "gene_id"
208 counts.attrs["provenance"] = {"source": "star", "strandedness": strandedness,
209 "rounding": "none", "inputs": [_file_record(f) for f in tabs]}
210 return counts
213def build_from_featurecounts(counts_file: Path) -> pd.DataFrame:
214 """Parse a combined featureCounts matrix into a gene x sample matrix."""
215 if not counts_file.is_file():
216 sys.exit(f"featureCounts file not found: {counts_file}")
217 # featureCounts prepends a '#' command line; real header is the next row.
218 validate_header(counts_file, delimiter="\t", comment="#")
219 df = pd.read_csv(counts_file, sep="\t", comment="#", dtype={"Geneid": str}, keep_default_na=False)
220 # Layout: Geneid, Chr, Start, End, Strand, Length, <bam1>, <bam2>, ...
221 meta_cols = ["Geneid", "Chr", "Start", "End", "Strand", "Length"]
222 if not set(meta_cols).issubset(df.columns):
223 raise ValueError("featureCounts file is missing required annotation columns")
224 sample_cols = [c for c in df.columns if c not in meta_cols]
225 if not sample_cols:
226 sys.exit("No sample/count columns found in featureCounts file")
227 counts = df.set_index("Geneid")[sample_cols]
228 counts.columns = [_clean_sample_name(c) for c in counts.columns]
229 counts = _validate_counts(counts)
230 counts.index.name = "gene_id"
231 counts.attrs["provenance"] = {"source": "featurecounts", "rounding": "none",
232 "inputs": [_file_record(counts_file)]}
233 print(f"featureCounts: {len(sample_cols)} samples -> {list(counts.columns)}")
234 return counts
237def write_outputs(counts: pd.DataFrame, output_dir: Path) -> None:
238 provenance = counts.attrs.get("provenance", {"source": "unspecified"})
239 counts = _validate_counts(counts)
240 if (counts.sum(axis=0) <= 0).any():
241 raise ValueError("A sample has zero total counts; resolve it before DE")
243 # Drop all-zero genes (uninformative; pydeseq2 filters further).
244 n_before = counts.shape[0]
245 counts = counts[counts.sum(axis=1) > 0]
246 dropped = n_before - counts.shape[0]
248 output_dir.mkdir(parents=True, exist_ok=True)
249 counts_path = output_dir / "counts.csv"
250 counts.to_csv(counts_path)
252 meta = pd.DataFrame(
253 {"condition": ["CHANGE_ME"] * counts.shape[1], "batch": [""] * counts.shape[1]},
254 index=pd.Index(counts.columns, name="sample"),
255 )
256 meta_path = output_dir / "metadata_template.csv"
257 meta.to_csv(meta_path)
258 provenance.update({"orientation": "genes_by_samples", "samples": list(counts.columns),
259 "all_zero_genes_removed": dropped, "pandas_version": pd.__version__})
260 (output_dir / "counts_provenance.json").write_text(json.dumps(provenance, indent=2) + "\n", encoding="utf-8")
262 print(f"\n genes: {counts.shape[0]} (dropped {dropped} all-zero)")
263 print(f" samples: {counts.shape[1]}")
264 print(f" wrote {counts_path} (genes x samples, integer)")
265 print(f" wrote {meta_path} (fill in 'condition'/'batch', then run the pydeseq2 skill)")
268def main() -> None:
269 p = argparse.ArgumentParser(
270 description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
271 )
272 p.add_argument(
273 "--from", dest="source", required=True,
274 choices=["salmon", "star", "featurecounts"],
275 help="Quantifier that produced the input.",
276 )
277 p.add_argument("--quant-dir", type=Path,
278 help="Directory of per-sample quant output (salmon/star).")
279 p.add_argument("--tx2gene", type=Path,
280 help="transcript_id->gene_id map (salmon). TSV/CSV with those columns.")
281 p.add_argument("--strandedness", choices=list(STAR_STRAND_COL),
282 help="Verified library strandedness (required for STAR; no default).")
283 p.add_argument("--ignore-transcript-version", action="store_true",
284 help="Strip terminal .N versions in Salmon IDs and map, after checking reference provenance.")
285 p.add_argument("--counts-file", type=Path,
286 help="Combined featureCounts matrix (featurecounts).")
287 p.add_argument("--output-dir", type=Path, default=Path("counts"),
288 help="Where to write counts.csv + metadata_template.csv (default: counts/).")
289 args = p.parse_args()
291 if args.source == "salmon":
292 if not args.quant_dir:
293 sys.exit("--quant-dir is required for --from salmon")
294 counts = build_from_salmon(args.quant_dir, args.tx2gene, args.ignore_transcript_version)
295 elif args.source == "star":
296 if not args.quant_dir:
297 sys.exit("--quant-dir is required for --from star")
298 if not args.strandedness:
299 sys.exit("--strandedness is required for --from star; infer it before selecting counts")
300 counts = build_from_star(args.quant_dir, args.strandedness)
301 else:
302 if not args.counts_file:
303 sys.exit("--counts-file is required for --from featurecounts")
304 counts = build_from_featurecounts(args.counts_file)
306 write_outputs(counts, args.output_dir)
309if __name__ == "__main__":
310 main()