Skills › Scientific Agent Skills › geniml › consensus_plan.py
consensus_plan.py
416 lines · skills/geniml/scripts/consensus_plan.py at 5f2bc03
1#!/usr/bin/env python3
2"""Create a bounded, local-only Geniml consensus-universe execution plan."""
4from __future__ import annotations
6import argparse
7import re
8import sys
9from collections import Counter
10from pathlib import Path
12from _common import (
13 HARD_MAX_BYTES,
14 HARD_MAX_FILES,
15 SafetyError,
16 add_path_mode_argument,
17 delimiter_from_path,
18 display_path,
19 fail_json,
20 int_type,
21 load_chrom_sizes,
22 local_path,
23 print_json,
24 read_delimited_manifest,
25 sha256_file,
26)
29TOOL = "consensus-peaks-planner"
30_SAFE_PREFIX = re.compile(r"^[A-Za-z0-9_.-]{1,100}$")
33def build_parser() -> argparse.ArgumentParser:
34 parser = argparse.ArgumentParser(
35 description=(
36 "Validate local consensus inputs and emit argv arrays for Geniml "
37 "0.8.4. No command, coverage tool, or network request is executed."
38 )
39 )
40 parser.add_argument("--manifest", required=True)
41 parser.add_argument("--chrom-sizes", required=True)
42 parser.add_argument("--assembly", required=True)
43 parser.add_argument("--output-dir", required=True)
44 parser.add_argument("--coverage-dir")
45 parser.add_argument("--likelihood-model")
46 parser.add_argument("--method", choices=("cc", "ccf", "ml", "hmm"), required=True)
47 parser.add_argument("--coverage-prefix", default="all")
48 parser.add_argument(
49 "--cutoff",
50 type=int_type(minimum=1, maximum=2**31 - 1, label="cutoff"),
51 )
52 parser.add_argument(
53 "--merge",
54 type=int_type(minimum=0, maximum=10_000_000, label="merge"),
55 default=0,
56 )
57 parser.add_argument(
58 "--filter-size",
59 type=int_type(minimum=0, maximum=10_000_000, label="filter-size"),
60 default=0,
61 )
62 parser.add_argument("--no-normalize", action="store_true")
63 parser.add_argument("--save-max-coverage", action="store_true")
64 parser.add_argument("--path-column", default="path")
65 parser.add_argument("--assembly-column", default="assembly")
66 parser.add_argument(
67 "--delimiter",
68 choices=("auto", "csv", "tsv"),
69 default="auto",
70 )
71 parser.add_argument(
72 "--max-files",
73 type=int_type(minimum=1, maximum=HARD_MAX_FILES, label="max-files"),
74 default=10_000,
75 )
76 parser.add_argument(
77 "--max-manifest-bytes",
78 type=int_type(
79 minimum=1,
80 maximum=HARD_MAX_BYTES,
81 label="max-manifest-bytes",
82 ),
83 default=64 * 1024**2,
84 )
85 parser.add_argument(
86 "--max-total-input-bytes",
87 type=int_type(
88 minimum=1,
89 maximum=HARD_MAX_BYTES,
90 label="max-total-input-bytes",
91 ),
92 default=4 * 1024**3,
93 )
94 add_path_mode_argument(parser)
95 return parser
98def _manifest_bed(raw: str, manifest: Path) -> Path:
99 candidate = Path(raw)
100 if not candidate.is_absolute():
101 candidate = manifest.parent / candidate
102 return local_path(str(candidate), kind="file")
105def _render(path: Path, label: str, mode: str) -> str:
106 if mode == "full":
107 return str(path)
108 if mode == "basename":
109 return path.name
110 return f"<{label}>"
113def plan(args: argparse.Namespace) -> tuple[dict, int]:
114 if not args.assembly.strip() or len(args.assembly) > 200:
115 raise SafetyError("assembly must be a nonempty value of at most 200 characters")
116 if not _SAFE_PREFIX.fullmatch(args.coverage_prefix):
117 raise SafetyError("coverage prefix contains unsupported characters")
118 if args.cutoff is not None and args.cutoff < 1:
119 raise SafetyError("cutoff must be at least 1")
120 if args.method != "cc" and (
121 args.cutoff is not None or args.merge != 0 or args.filter_size != 0
122 ):
123 raise SafetyError("cutoff, merge, and filter-size apply only to method cc")
124 if args.method != "hmm" and (args.no_normalize or args.save_max_coverage):
125 raise SafetyError("normalization and max-coverage flags apply only to method hmm")
126 if args.likelihood_model and args.method != "ml":
127 raise SafetyError("likelihood-model applies only to method ml")
129 manifest = local_path(args.manifest, kind="file")
130 chrom_sizes = local_path(args.chrom_sizes, kind="file")
131 output_dir = local_path(args.output_dir, kind="dir")
132 delimiter = delimiter_from_path(manifest, args.delimiter)
133 fields, rows = read_delimited_manifest(
134 manifest,
135 delimiter_name=delimiter,
136 max_bytes=args.max_manifest_bytes,
137 max_rows=args.max_files,
138 )
139 if args.path_column not in fields:
140 raise SafetyError(f"path column is missing: {args.path_column}")
141 if args.assembly_column not in fields:
142 raise SafetyError(f"assembly column is missing: {args.assembly_column}")
143 if not rows:
144 raise SafetyError("manifest has no data rows")
146 sizes, contig_order = load_chrom_sizes(
147 chrom_sizes,
148 max_bytes=min(args.max_manifest_bytes, 128 * 1024**2),
149 max_records=1_000_000,
150 )
151 manifest_digest, manifest_size = sha256_file(
152 manifest,
153 max_bytes=args.max_manifest_bytes,
154 )
155 chrom_digest, chrom_size_bytes = sha256_file(
156 chrom_sizes,
157 max_bytes=min(args.max_manifest_bytes, 128 * 1024**2),
158 )
160 errors: Counter[str] = Counter()
161 warnings: Counter[str] = Counter()
162 local_beds: list[Path] = []
163 total_input_bytes = 0
164 for row in rows:
165 raw_path = row.get(args.path_column, "")
166 if not raw_path:
167 errors["blank_path"] += 1
168 continue
169 if row.get(args.assembly_column, "") != args.assembly:
170 errors["assembly_mismatch"] += 1
171 try:
172 bed = _manifest_bed(raw_path, manifest)
173 except (OSError, SafetyError):
174 errors["invalid_or_missing_local_bed"] += 1
175 continue
176 local_beds.append(bed)
177 total_input_bytes += bed.lstat().st_size
178 if total_input_bytes > args.max_total_input_bytes:
179 raise SafetyError("BED corpus exceeds max-total-input-bytes")
180 if len(set(local_beds)) != len(local_beds):
181 errors["duplicate_bed_paths"] = len(local_beds) - len(set(local_beds))
183 coverage_dir: Path | None = None
184 required_coverage_names = (
185 [f"{args.coverage_prefix}_core.bw"]
186 if args.method in {"cc", "ccf"}
187 else [
188 f"{args.coverage_prefix}_start.bw",
189 f"{args.coverage_prefix}_core.bw",
190 f"{args.coverage_prefix}_end.bw",
191 ]
192 )
193 coverage_files: list[dict[str, str | int]] = []
194 if args.coverage_dir:
195 coverage_dir = local_path(args.coverage_dir, kind="dir")
196 for index, name in enumerate(required_coverage_names, start=1):
197 try:
198 coverage_file = local_path(str(coverage_dir / name), kind="file")
199 except (OSError, SafetyError):
200 errors[f"missing_coverage:{name}"] += 1
201 continue
202 digest, size = sha256_file(
203 coverage_file,
204 max_bytes=args.max_total_input_bytes,
205 )
206 coverage_files.append(
207 {
208 "name": name,
209 "sha256": digest,
210 "size_bytes": size,
211 }
212 )
213 else:
214 warnings["coverage_generation_required"] += 1
216 output_file = output_dir / f"universe_{args.method}.bed"
217 if output_file.exists() or output_file.is_symlink():
218 errors["output_file_already_exists"] += 1
220 coverage_arg = (
221 _render(coverage_dir, "LOCAL_COVERAGE_DIR", args.path_mode)
222 if coverage_dir
223 else "<LOCAL_COVERAGE_DIR>"
224 )
225 output_arg = _render(output_file, "OUTPUT_UNIVERSE_BED", args.path_mode)
227 stages: list[dict] = [
228 {
229 "stage": "validate_beds",
230 "executed": False,
231 "instruction": (
232 "Run bed_validator.py for every BED using the same chromosome sizes; "
233 "reject any invalid result."
234 ),
235 },
236 {
237 "stage": "generate_coverage",
238 "executed": False,
239 "blocked": coverage_dir is None,
240 "required_outputs": required_coverage_names,
241 "instruction": (
242 "Use an explicitly pinned and checksummed coverage executable. "
243 "Capture its installed --help; this planner does not guess a Gtars/uniwig argv."
244 ),
245 },
246 ]
248 likelihood_path: Path | None = None
249 if args.method == "ml":
250 if args.likelihood_model:
251 likelihood_path = local_path(args.likelihood_model, kind="file")
252 else:
253 likelihood_path = output_dir / "likelihood_model.tar"
254 if likelihood_path.exists() or likelihood_path.is_symlink():
255 errors["planned_likelihood_model_already_exists"] += 1
256 lh_template = [
257 "geniml",
258 "lh",
259 "--model-file",
260 "<LIKELIHOOD_MODEL_TAR>",
261 "--coverage-folder",
262 "<LOCAL_COVERAGE_DIR>",
263 "--coverage-prefix",
264 args.coverage_prefix,
265 "--file-no",
266 str(len(set(local_beds))),
267 ]
268 lh_argv = None
269 if args.path_mode == "full" and coverage_dir:
270 lh_argv = [
271 "geniml",
272 "lh",
273 "--model-file",
274 str(likelihood_path),
275 "--coverage-folder",
276 str(coverage_dir),
277 "--coverage-prefix",
278 args.coverage_prefix,
279 "--file-no",
280 str(len(set(local_beds))),
281 ]
282 stages.append(
283 {
284 "stage": "build_likelihood_model",
285 "executed": False,
286 "argv_template": lh_template,
287 "argv": lh_argv,
288 }
289 )
291 command_template = [
292 "geniml",
293 "build-universe",
294 args.method,
295 "--coverage-folder",
296 "<LOCAL_COVERAGE_DIR>",
297 "--coverage-prefix",
298 args.coverage_prefix,
299 "--output-file",
300 "<OUTPUT_UNIVERSE_BED>",
301 ]
302 command = [
303 "geniml",
304 "build-universe",
305 args.method,
306 "--coverage-folder",
307 coverage_arg,
308 "--coverage-prefix",
309 args.coverage_prefix,
310 "--output-file",
311 output_arg,
312 ]
313 if args.method == "cc":
314 if args.cutoff is not None:
315 command_template.extend(["--cutoff", str(args.cutoff)])
316 command.extend(["--cutoff", str(args.cutoff)])
317 if args.cutoff > len(set(local_beds)):
318 warnings["cutoff_exceeds_file_count"] += 1
319 command_template.extend(["--merge", str(args.merge), "--filter-size", str(args.filter_size)])
320 command.extend(["--merge", str(args.merge), "--filter-size", str(args.filter_size)])
321 elif args.method == "ml":
322 command_template.extend(["--model-file", "<LIKELIHOOD_MODEL_TAR>"])
323 command.extend(
324 [
325 "--model-file",
326 _render(
327 likelihood_path,
328 "LIKELIHOOD_MODEL_TAR",
329 args.path_mode,
330 ),
331 ]
332 )
333 elif args.method == "hmm":
334 if args.no_normalize:
335 command_template.append("--not-normalize")
336 command.append("--not-normalize")
337 if args.save_max_coverage:
338 command_template.append("--save-max-cove")
339 command.append("--save-max-cove")
341 stages.append(
342 {
343 "stage": "build_universe",
344 "executed": False,
345 "argv_template": command_template,
346 "argv": command if args.path_mode == "full" and coverage_dir else None,
347 }
348 )
349 stages.append(
350 {
351 "stage": "validate_output",
352 "executed": False,
353 "instruction": (
354 "Validate BED coordinates, sorting, contigs, bounds, nonempty output, "
355 "and checksum before constructing a tokenizer."
356 ),
357 }
358 )
360 ready = not errors and coverage_dir is not None
361 report = {
362 "ok": not errors,
363 "ready_to_execute": ready,
364 "tool": TOOL,
365 "contract": {
366 "assembly": args.assembly,
367 "coordinate_system": "0-based-half-open",
368 "method": args.method,
369 "network_used": False,
370 "commands_executed": False,
371 "paths_disclosed": args.path_mode == "full",
372 },
373 "inputs": {
374 "manifest": {
375 "path": display_path(manifest, 1, args.path_mode),
376 "sha256": manifest_digest,
377 "size_bytes": manifest_size,
378 },
379 "chromosome_sizes": {
380 "path": display_path(chrom_sizes, 2, args.path_mode),
381 "sha256": chrom_digest,
382 "size_bytes": chrom_size_bytes,
383 "contig_count": len(sizes),
384 "first_contigs": contig_order[:10],
385 },
386 "bed_rows": len(rows),
387 "unique_local_beds": len(set(local_beds)),
388 "total_bed_bytes": total_input_bytes,
389 "coverage_files": coverage_files,
390 },
391 "errors": dict(sorted(errors.items())),
392 "warnings": dict(sorted(warnings.items())),
393 "stages": stages,
394 "provenance_requirements": [
395 "patient/donor-grouped training-only manifest",
396 "input BED and chromosome-sizes checksums",
397 "coverage executable version, checksum, help output, and argv",
398 "coverage-track and universe checksums",
399 "Geniml/Gtars/Python lockfile and platform",
400 ],
401 }
402 return report, 0 if not errors else 2
405def main() -> int:
406 args = build_parser().parse_args()
407 try:
408 report, exit_code = plan(args)
409 print_json(report)
410 return exit_code
411 except (OSError, SafetyError, UnicodeError) as exc:
412 return fail_json(TOOL, exc)
415if __name__ == "__main__":
416 sys.exit(main())