Skills › Scientific Agent Skills › mageck › screen_analysis.py
screen_analysis.py
259 lines · skills/mageck/scripts/screen_analysis.py at 5f2bc03
1"""Validate pooled-screen inputs, quantify QC, and execute MAGeCK RRA testing."""
2from __future__ import annotations
4import argparse
5import csv
6from collections import Counter
7import hashlib
8import itertools
9import json
10import math
11from pathlib import Path
12import re
13import shutil
14import statistics
15import subprocess
16import tempfile
19def numeric_label(value):
20 """Upstream guesses numeric headers as counts and integer labels as indices."""
21 try:
22 float(value)
23 return True
24 except ValueError:
25 return False
28def validate_identifier(value):
29 if not value or any(c.isspace() for c in value):
30 raise ValueError("Guide and gene IDs must be nonempty and free of whitespace")
33def read_counts(path):
34 with Path(path).open(newline="") as handle:
35 rows = list(csv.reader(handle, delimiter="\t"))
36 if len(rows) < 2 or len(rows[0]) < 4:
37 raise ValueError("Expected sgRNA, Gene, and at least two sample columns")
38 if [x.lower() for x in rows[0][:2]] != ["sgrna", "gene"]:
39 raise ValueError("Count table requires a header beginning sgRNA and Gene")
40 samples = rows[0][2:]
41 if len(set(samples)) != len(samples) or any(not x or ',' in x or any(c.isspace() for c in x) for x in samples):
42 raise ValueError("Sample labels must be unique, nonempty, and free of commas/whitespace")
43 if any(numeric_label(x) for x in samples):
44 raise ValueError("Numeric sample labels are ambiguous in MAGeCK; rename with a prefix such as sample_1")
45 guides, genes, counts = [], [], []
46 for line, row in enumerate(rows[1:], 2):
47 if len(row) != len(rows[0]) or not row[0] or not row[1]:
48 raise ValueError(f"Invalid row {line}: missing fields")
49 validate_identifier(row[0]); validate_identifier(row[1])
50 try:
51 values = [int(v) for v in row[2:]]
52 except ValueError as error:
53 raise ValueError(f"Row {line}: counts must be integers") from error
54 if min(values) < 0:
55 raise ValueError(f"Row {line}: negative count")
56 guides.append(row[0]); genes.append(row[1]); counts.append(values)
57 if len(set(guides)) != len(guides):
58 raise ValueError("Duplicate guide IDs")
59 if any(sum(col) == 0 for col in zip(*counts)):
60 raise ValueError("A sample has zero total reads")
61 return samples, guides, genes, counts
64def read_library(path):
65 """MAGeCK headerless tab-separated guide, sequence, gene library."""
66 library, sequences = {}, set()
67 with Path(path).open(newline="") as handle:
68 for line, row in enumerate(csv.reader(handle, delimiter="\t"), 1):
69 if len(row) != 3 or not all(row):
70 raise ValueError(f"Library row {line}: expected guide, sequence, gene (no header)")
71 guide, seq, gene = row
72 validate_identifier(guide); validate_identifier(gene)
73 seq = seq.upper()
74 if set(seq) - set("ACGT") or guide in library or seq in sequences:
75 raise ValueError(f"Library row {line}: invalid DNA, duplicate ID, or ambiguous sequence")
76 library[guide] = (seq, gene); sequences.add(seq)
77 if not library:
78 raise ValueError("Empty library")
79 return library
82def gini(values):
83 ordered = sorted(values)
84 return 2 * sum(i * x for i, x in enumerate(ordered, 1)) / (len(ordered) * sum(ordered)) - (len(ordered) + 1) / len(ordered)
87def correlation(a, b):
88 a = [math.log2(x + 1) for x in a]; b = [math.log2(x + 1) for x in b]
89 mean_a, mean_b = statistics.mean(a), statistics.mean(b)
90 ac = [x - mean_a for x in a]; bc = [x - mean_b for x in b]
91 denominator = math.sqrt(sum(x*x for x in ac) * sum(x*x for x in bc))
92 return sum(x*y for x, y in zip(ac, bc)) / denominator if denominator else None
95def inspect_screen(count_path, control, treatment, library_path=None, controls_path=None, paired=False):
96 samples, guides, genes, counts = read_counts(count_path)
97 if not control or not treatment or len(set(control + treatment)) != len(control + treatment):
98 raise ValueError("Control and treatment must be nonempty, disjoint lists without duplicates")
99 if not set(control + treatment) <= set(samples):
100 raise ValueError("Unknown sample label in contrast")
101 if paired and len(control) != len(treatment):
102 raise ValueError("Paired contrasts need equal counts in matching biological order")
103 if library_path:
104 library = read_library(library_path)
105 if set(library) != set(guides):
106 raise ValueError("Count/library guide sets differ; retain zero-count library guides")
107 if any(library[g][1] != gene for g, gene in zip(guides, genes)):
108 raise ValueError("Count/library gene annotations differ")
109 control_guides = []
110 if controls_path:
111 control_guides = Path(controls_path).read_text().splitlines()
112 if not control_guides or len(set(control_guides)) != len(control_guides) or not set(control_guides) <= set(guides):
113 raise ValueError("Control-guide file must contain unique known IDs, one per line")
114 if len(control_guides) < 2:
115 raise ValueError("MAGeCK requires at least two control-guide IDs")
116 control_set = set(control_guides)
117 control_genes = {gene for guide, gene in zip(guides, genes) if guide in control_set}
118 if any(gene in control_genes and guide not in control_set for guide, gene in zip(guides, genes)):
119 raise ValueError("A gene cannot mix control and non-control guides in MAGeCK")
120 control_rows = [row for guide, row in zip(guides, counts) if guide in control_set]
121 if any(sum(col) == 0 for col in zip(*control_rows)):
122 raise ValueError("Control-guide counts are zero in a sample")
123 columns = dict(zip(samples, zip(*counts)))
124 qc = {s: {"reads": sum(v), "zero_fraction": v.count(0)/len(v),
125 "median_reads_per_guide": statistics.median(v), "gini": gini(v)} for s, v in columns.items()}
126 replicate_correlations = {f"{a} vs {b}": correlation(columns[a], columns[b])
127 for group in (control, treatment) for a, b in itertools.combinations(group, 2)}
128 guide_multiplicity = dict(Counter(genes))
129 warnings = []
130 if min(len(control), len(treatment)) < 2:
131 warnings.append("A condition lacks biological replication; dispersion and reproducibility are limited")
132 if any(v["zero_fraction"] > .1 for v in qc.values()):
133 warnings.append("More than 10% zero-count guides in a sample: review bottlenecks and mapping")
134 if any(v is not None and v < .8 for v in replicate_correlations.values()):
135 warnings.append("Low within-condition log-count correlation (<0.8): inspect replicate effects")
136 if any(v is None for v in replicate_correlations.values()):
137 warnings.append("Some replicate correlations are undefined because a sample has constant counts")
138 if min(guide_multiplicity.values()) < 3:
139 warnings.append("Some genes have fewer than three guides; inspect guide concordance")
140 return {"samples": qc, "guides": len(guides), "genes": len(guide_multiplicity),
141 "guides_per_gene": guide_multiplicity, "within_condition_log2_count_correlations": replicate_correlations,
142 "control": control, "treatment": treatment, "paired": paired,
143 "control_guides": len(control_guides), "warnings": warnings}
146def normalization_evidence(log, requested, samples):
147 """Extract applied scaling from the 0.5.9.5 log, including sparse-count fallback."""
148 matches = re.findall(r"Final size factor: ([^\r\n]+)", log)
149 factors = [float(x) for x in matches[-1].split()] if matches else []
150 if len(factors) != len(samples) or any(not math.isfinite(x) or x <= 0 for x in factors):
151 raise RuntimeError("Missing or invalid MAGeCK normalization factors; inspect execution.log")
152 fallback = "Switch to total read count normalization." in log
153 return {"requested_method": requested,
154 "effective_method": "total" if fallback else "median" if requested == "control" else requested,
155 "reference_guides": "control_guides" if requested == "control" else "all_guides",
156 "fallback_to_total": fallback,
157 "sample_size_factors": dict(zip(samples, factors)),
158 "warnings": [line for line in log.splitlines() if "WARNING" in line]}
161def read_gene_hits(path, fdr):
162 hits = {"neg": [], "pos": []}
163 seen = set()
164 with Path(path).open() as handle:
165 rows = csv.DictReader(handle, delimiter="\t")
166 for row in rows:
167 gene = row["id"]
168 if not gene or gene in seen:
169 raise ValueError("Missing or duplicate gene in MAGeCK summary")
170 seen.add(gene)
171 for direction in hits:
172 qvalue = float(row[f"{direction}|fdr"])
173 lfc = float(row[f"{direction}|lfc"])
174 rank = int(row[f"{direction}|rank"])
175 if not math.isfinite(qvalue) or not 0 <= qvalue <= 1 or not math.isfinite(lfc) or rank < 1:
176 raise ValueError("Invalid FDR, effect size, or rank in MAGeCK summary")
177 if qvalue <= fdr:
178 hits[direction].append({"gene": gene, "fdr": qvalue, "lfc": lfc, "rank": rank})
179 if not seen:
180 raise ValueError("Empty MAGeCK gene summary")
181 return hits
184def run_test(count_path, output, control, treatment, library_path=None, controls_path=None,
185 paired=False, normalization="median", executable="mageck", fdr=.05):
186 if not 0 < fdr < 1:
187 raise ValueError("FDR must lie between zero and one")
188 if normalization not in {"median", "total", "control", "none"}:
189 raise ValueError("Unsupported normalization")
190 if normalization == "control" and not controls_path:
191 raise ValueError("Control normalization requires control-guide IDs")
192 qc = inspect_screen(count_path, control, treatment, library_path, controls_path, paired)
193 binary = shutil.which(executable)
194 if binary is None:
195 raise RuntimeError("MAGeCK executable not found; install MAGeCK 0.5.9.5 and its RRA binary")
196 output = Path(output)
197 output.mkdir(parents=True, exist_ok=False)
198 version = subprocess.run([binary, "--version"], text=True, capture_output=True, check=True).stdout.strip()
199 # Upstream uses shell commands internally. Fixed staging names avoid its path quoting problems.
200 with tempfile.TemporaryDirectory(prefix="mageck-analysis-") as temporary:
201 work = Path(temporary)
202 shutil.copyfile(count_path, work / "counts.tsv")
203 command = [binary, "test", "-k", "counts.tsv", "-c", ",".join(control), "-t", ",".join(treatment),
204 "--norm-method", normalization, "--normcounts-to-file", "--gene-lfc-method", "median",
205 "--remove-zero", "both", "-n", "screen"]
206 if paired:
207 command.append("--paired")
208 if controls_path:
209 shutil.copyfile(controls_path, work / "controls.txt")
210 command += ["--control-sgrna", "controls.txt"]
211 process = subprocess.run(command, cwd=work, capture_output=True, text=True, check=False)
212 log = process.stdout + process.stderr
213 (output / "execution.log").write_text(log)
214 for artifact in work.glob("screen.*"):
215 if artifact.is_file():
216 shutil.copyfile(artifact, output / artifact.name)
217 if process.returncode or not (output / "screen.gene_summary.txt").is_file():
218 raise RuntimeError(f"MAGeCK did not complete; inspect {output / 'execution.log'}")
219 # 0.5.9.5 normalizes cttab_sel in selected_samples_id = controlgroup + treatgroup
220 # order before any paired-guide expansion, rather than in original count-column order.
221 normalization_info = normalization_evidence(log, normalization, control + treatment)
222 hits = read_gene_hits(output / "screen.gene_summary.txt", fdr)
223 report = {"mageck_version": version, "command": command, "count_sha256": hashlib.sha256(Path(count_path).read_bytes()).hexdigest(),
224 "library_sha256": hashlib.sha256(Path(library_path).read_bytes()).hexdigest() if library_path else None,
225 "control_guides_sha256": hashlib.sha256(Path(controls_path).read_bytes()).hexdigest() if controls_path else None,
226 "control_guide_ids": Path(controls_path).read_text().splitlines() if controls_path else [],
227 "normalization": normalization_info, "fdr_threshold": fdr, "qc": qc, "hits": hits}
228 (output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
229 return report
232def main():
233 parser = argparse.ArgumentParser(description=__doc__)
234 parser.add_argument("mode", choices=["qc", "test"])
235 parser.add_argument("--counts", required=True, type=Path)
236 parser.add_argument("--control", nargs="+", required=True)
237 parser.add_argument("--treatment", nargs="+", required=True)
238 parser.add_argument("--library", type=Path)
239 parser.add_argument("--control-guides", type=Path)
240 parser.add_argument("--paired", action="store_true")
241 parser.add_argument("--normalization", choices=["median", "total", "control", "none"], default="median")
242 parser.add_argument("--output", type=Path, help="New directory for test output; JSON path for QC")
243 parser.add_argument("--mageck", default="mageck")
244 parser.add_argument("--fdr", type=float, default=.05)
245 args = parser.parse_args()
246 if args.mode == "test":
247 if args.output is None:
248 parser.error("test requires --output")
249 report = run_test(args.counts, args.output, args.control, args.treatment, args.library,
250 args.control_guides, args.paired, args.normalization, args.mageck, args.fdr)
251 else:
252 report = inspect_screen(args.counts, args.control, args.treatment, args.library, args.control_guides, args.paired)
253 if args.output:
254 args.output.write_text(json.dumps(report, indent=2) + "\n")
255 print(json.dumps(report, indent=2))
258if __name__ == "__main__":
259 main()