Spaces:
Sleeping
Sleeping
| """Simulator orchestrator: assign each pair a label, rewrite R2 for affected pairs (R1 copied | |
| byte-identically), and emit the failure FASTQs, ground-truth labels, and a run manifest.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import platform | |
| import shutil | |
| from collections import Counter | |
| from pathlib import Path | |
| import numpy as np | |
| from seqcolyte.io.fastx import FastqRecord, iter_pairs, write_fastx_gz | |
| from seqcolyte.spec.loader import load_spec | |
| from sim import SIM_VERSION | |
| from sim.base import ReadCtx | |
| from sim.config import SimConfig | |
| from sim.labels import write_labels | |
| from sim.registry import get_mode | |
| from sim.rng import read_rng | |
| __all__ = ["run_simulation", "assign_label"] | |
| def assign_label(rng: np.random.Generator, params: dict) -> tuple[str, str | None]: | |
| """Assign clean / readthrough / pure_dimer using two sequential draws from ``rng``.""" | |
| affected_fraction = float(params.get("affected_fraction", 0.30)) | |
| dimer_fraction = float(params.get("dimer_fraction", 0.33)) | |
| if rng.random() >= affected_fraction: | |
| return "clean", None | |
| if rng.random() < dimer_fraction: | |
| return "pure_dimer", "pure_dimer" | |
| return "readthrough", "readthrough" | |
| def _md5(path: str) -> str: | |
| h = hashlib.md5() | |
| with open(path, "rb") as fh: | |
| for chunk in iter(lambda: fh.read(1 << 20), b""): | |
| h.update(chunk) | |
| return h.hexdigest() | |
| def run_simulation(config: SimConfig) -> dict: | |
| spec = load_spec(config.spec) | |
| mode = get_mode(config.failure_mode) | |
| if not mode.applies_to(spec): | |
| raise ValueError( | |
| f"mode {mode.name!r} (platform {mode.platform!r}) does not apply to spec platform {spec.platform!r}" | |
| ) | |
| cb_slice = spec.segment_slice("R1", "cell_barcode") | |
| umi_slice = spec.segment_slice("R1", "umi") | |
| for p in (config.out_r1, config.out_r2, config.out_labels, config.out_manifest): | |
| Path(p).parent.mkdir(parents=True, exist_ok=True) | |
| # R1 is left byte-identical: copy the file rather than round-tripping records. | |
| shutil.copyfile(config.input_r1, config.out_r1) | |
| counts: Counter[str] = Counter() | |
| rows: list[dict] = [] | |
| def rewritten_r2(): | |
| for i, (r1, r2) in enumerate(iter_pairs(config.input_r1, config.input_r2)): | |
| rng = read_rng(config.seed, i) | |
| label, subtype = assign_label(rng, config.params) | |
| cb = r1.sequence[cb_slice] | |
| umi = r1.sequence[umi_slice] | |
| if label == "clean": | |
| out = r2 | |
| fields: dict = {} | |
| construct = "clean" | |
| else: | |
| ctx = ReadCtx(r1=r1, r2=r2, spec=spec, params=config.params, rng=rng, | |
| pair_index=i, subtype=subtype, cb=cb, umi=umi, r2_len=len(r2.sequence)) | |
| res = mode.build_r2(ctx) | |
| out = FastqRecord(name=r2.name, comment=r2.comment, sequence=res.sequence, quality=res.quality) | |
| fields = res.fields | |
| construct = res.construct | |
| counts[label] += 1 | |
| rows.append({ | |
| "read_id": r2.name, "pair_index": i, "label": label, | |
| "failure_mode": config.failure_mode if label != "clean" else "none", | |
| "affected": int(label != "clean"), | |
| "r1_len": len(r1.sequence), "r2_len": len(out.sequence), | |
| "cb": cb, "umi": umi, | |
| "insert_len": fields.get("insert_len", ""), | |
| "polyA_len": fields.get("polyA_len", ""), | |
| "pad_len": fields.get("pad_len", ""), | |
| "truncated": int(fields["truncated"]) if "truncated" in fields else "", | |
| "construct": construct, | |
| }) | |
| yield out | |
| n_pairs = write_fastx_gz(config.out_r2, rewritten_r2()) | |
| write_labels(config.out_labels, rows) | |
| r1_in_md5, r1_out_md5 = _md5(config.input_r1), _md5(config.out_r1) | |
| r1_identical = r1_in_md5 == r1_out_md5 | |
| manifest = { | |
| "name": config.name, | |
| "sim_version": SIM_VERSION, | |
| "spec_id": spec.spec_id, | |
| "failure_mode": config.failure_mode, | |
| "seed": config.seed, | |
| "params": config.params, | |
| "n_pairs": n_pairs, | |
| "label_counts": dict(counts), | |
| "label_fractions": {k: round(v / n_pairs, 4) for k, v in counts.items()} if n_pairs else {}, | |
| "inputs": {"r1": config.input_r1, "r2": config.input_r2}, | |
| "outputs": {"r1": config.out_r1, "r2": config.out_r2, "labels": config.out_labels}, | |
| "r1_input_md5": r1_in_md5, | |
| "r1_output_md5": r1_out_md5, | |
| "r1_byte_identical": r1_identical, | |
| "python": platform.python_version(), | |
| } | |
| Path(config.out_manifest).write_text(json.dumps(manifest, indent=2) + "\n") | |
| if not r1_identical: | |
| raise RuntimeError("R1 output is not byte-identical to input — invariant violated") | |
| return manifest | |