giant/scripts/create_root_files.py

"""Generate new ROOT shards by running a minicalosim executable (e.g.
run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree:

    raw/<kind>/<gen>/<detector>/shard-NNN.root

These executables take `[configName] nEvents [energy_GeV]` (configName is
only accepted by executables with a config selector, e.g. run_sampling;
energy_GeV defaults to 1.0 in the executable itself if omitted here) and
always write a fixed-name *.root file into the current directory — so
running several in parallel needs separate working directories, and the
output filename has to be discovered rather than assumed (it differs per
executable: run_pbwo4 writes pbwo4_<n>events_hits.root, run_sampling writes
sampling_<config>_<n>events_hits.root, others may differ again). This script
gives each run its own scratch directory under <dataset-root>/.sim-tmp/,
requires exactly one *.root to appear there, and moves it to the next free
shard index for that detector (existing shards are never overwritten).

--gen must already exist under raw/<kind>/ — create one first with
`dwarf bump-gen`.

See `uv run dwarf make-root --help` for the CLI.
"""

import os
import re
import shutil
import subprocess
import sys
import uuid
import zlib
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

# Must match scripts/bump_dataset_version.py's GEN_RE.
GEN_RE = re.compile(r"^gen\d+$")
SHARD_RE = re.compile(r"^shard-(\d+)\.root$")


class PlanError(ValueError):
    pass


@dataclass(frozen=True)
class SimJob:
    detector: str
    config: str | None
    shard_index: int


@dataclass(frozen=True)
class JobResult:
    job: SimJob
    ok: bool
    dest: Path | None
    message: str
    stdout: str
    stderr: str


def parse_detector_spec(spec: str) -> tuple[str, str | None]:
    """'sampling_pb_scint:pb_scint' -> ('sampling_pb_scint', 'pb_scint');
    'pbwo4' -> ('pbwo4', None)."""
    if ":" in spec:
        label, config = spec.split(":", 1)
        if not label or not config:
            raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
        return label, config
    return spec, None


def next_shard_index(detector_dir: Path) -> int:
    """One past the highest existing shard-NNN.root in *detector_dir* (0 if none/missing)."""
    if not detector_dir.is_dir():
        return 0
    best = -1
    for child in detector_dir.iterdir():
        m = SHARD_RE.match(child.name)
        if m and child.is_file():
            best = max(best, int(m.group(1)))
    return best + 1


def plan_jobs(
    detector_specs: list[str],
    num_files: int,
    dataset_root: Path,
    kind: str,
    gen: str,
) -> list[SimJob]:
    if not GEN_RE.match(gen):
        raise PlanError(f"--gen must look like 'genN', got {gen!r}")
    gen_dir = dataset_root / "raw" / kind / gen
    if not gen_dir.is_dir():
        raise PlanError(f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first")

    jobs = []
    for spec in detector_specs:
        label, config = parse_detector_spec(spec)
        start = next_shard_index(gen_dir / label)
        for i in range(num_files):
            jobs.append(SimJob(detector=label, config=config, shard_index=start + i))
    return jobs


def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
    """Deterministic RNG seed for one sim job, unique per (kind, gen, detector, config, shard, energy).

    Jobs run concurrently (ThreadPoolExecutor below) and can start within the
    same wall-clock second; minicalosim's default seed falls back to
    time(NULL) in that case, so two concurrently-launched jobs can silently
    get identical RNG state and produce byte-identical physics despite
    landing in separate shard files. Deriving the seed from the full job
    identity instead keeps it both unique and reproducible.
    """
    key = (
        f"{kind}|{gen}|{job.detector}|{job.config or ''}|{job.shard_index}"
        f"|{energy_gev if energy_gev is not None else ''}"
    )
    return zlib.crc32(key.encode()) & 0x7FFFFFFF


def build_cmd(executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None) -> list[str]:
    """minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
    cmd = [str(executable)]
    if job.config:
        cmd.append(job.config)
    cmd.append(str(events_per_file))
    if energy_gev is not None:
        cmd.append(str(energy_gev))
    return cmd


def run_job(
    job: SimJob,
    executable: Path,
    events_per_file: int,
    energy_gev: float | None,
    dataset_root: Path,
    kind: str,
    gen: str,
    tmp_root: Path,
) -> JobResult:
    workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
    workdir.mkdir(parents=True)

    cmd = build_cmd(executable, job, events_per_file, energy_gev)

    env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job, energy_gev)))
    result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env)

    if result.returncode != 0:
        return JobResult(
            job,
            False,
            None,
            f"executable exited {result.returncode}",
            result.stdout,
            result.stderr,
        )

    produced = sorted(workdir.glob("*.root"))
    if len(produced) != 1:
        return JobResult(
            job,
            False,
            None,
            f"expected exactly one .root output in {workdir}, found {len(produced)}: {[p.name for p in produced]}",
            result.stdout,
            result.stderr,
        )

    dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
    if dest.exists():
        return JobResult(
            job,
            False,
            None,
            f"refusing to overwrite existing {dest}",
            result.stdout,
            result.stderr,
        )

    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.move(str(produced[0]), str(dest))
    # rmtree, not rmdir: the executable may leave other side-effect files
    # (logs, seed state, ...) behind in its scratch workdir besides the .root.
    shutil.rmtree(workdir)
    return JobResult(job, True, dest, "ok", result.stdout, result.stderr)


def run_all(
    jobs: list[SimJob],
    executable: Path,
    events_per_file: int,
    energy_gev: float | None,
    dataset_root: Path,
    kind: str,
    gen: str,
    max_workers: int,
    tmp_root: Path,
) -> list[JobResult]:
    results = []
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(
                run_job,
                job,
                executable,
                events_per_file,
                energy_gev,
                dataset_root,
                kind,
                gen,
                tmp_root,
            ): job
            for job in jobs
        }
        for future in as_completed(futures):
            result = future.result()
            label = f"{result.job.detector} shard-{result.job.shard_index:03d}"
            status = "ok" if result.ok else f"FAILED: {result.message}"
            print(f"\n=== {label}: {status} ===")
            if result.stdout:
                print(result.stdout, end="")
            if result.stderr:
                print(result.stderr, end="", file=sys.stderr)
            results.append(result)
    return results


def run_make_root(
    executable: Path,
    detector: list[str],
    num_files: int,
    events_per_file: int,
    kind: str,
    gen: str,
    dataset_root: str,
    jobs: int,
    execute: bool,
    energy_gev: float | None = None,
) -> None:
    if jobs < 1:
        raise SystemExit("error: --jobs must be >= 1")
    if num_files < 1:
        raise SystemExit("error: --num-files must be >= 1")
    if events_per_file < 1:
        raise SystemExit("error: --events-per-file must be >= 1")
    if energy_gev is not None and energy_gev <= 0:
        raise SystemExit("error: --energy-gev must be > 0")
    if not executable.is_file() or not os.access(executable, os.X_OK):
        raise SystemExit(f"error: {executable} is not an executable file")

    dataset_root_path = Path(dataset_root)
    try:
        planned_jobs = plan_jobs(detector, num_files, dataset_root_path, kind, gen)
    except PlanError as exc:
        raise SystemExit(f"error: {exc}")

    print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
    print(f"executable: {executable}")
    for job in planned_jobs:
        cmd = build_cmd(executable, job, events_per_file, energy_gev)
        dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
        seed = job_seed(kind, gen, job, energy_gev)
        print(f"  MINICALOSIM_SEED={seed} {' '.join(cmd)}  ->  {dest}")

    if not execute:
        print("\nDry run only — pass --execute to apply.")
        return

    tmp_root = dataset_root_path / ".sim-tmp"
    tmp_root.mkdir(parents=True, exist_ok=True)
    results = run_all(
        planned_jobs,
        executable,
        events_per_file,
        energy_gev,
        dataset_root_path,
        kind,
        gen,
        max_workers=jobs,
        tmp_root=tmp_root,
    )
    if tmp_root.is_dir() and not any(tmp_root.iterdir()):
        tmp_root.rmdir()

    failures = [r for r in results if not r.ok]
    if failures:
        print(f"\n{len(failures)} of {len(results)} job(s) failed:", file=sys.stderr)
        for r in failures:
            print(
                f"  {r.job.detector} shard-{r.job.shard_index:03d}: {r.message}",
                file=sys.stderr,
            )
        raise SystemExit(1)

    print(f"\nAll {len(results)} job(s) completed.")