From c81076c324d8f44b5cc06a992efbfad35a29d5af Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 13 Jul 2026 13:59:09 +0200 Subject: [PATCH] Support --energy-gev in dwarf make-root for the new minicalosim energy arg run_pbwo4/run_sampling now accept a trailing energy_GeV positional arg; thread it through plan/run/seed so datasets like pbwo4_10gev can be generated at non-default beam energies. Co-Authored-By: Claude Sonnet 5 --- scripts/create_root_files.py | 64 +++++++++++++++++++++------------ scripts/dwarf.py | 11 ++++++ tests/test_create_root_files.py | 55 ++++++++++++++++++++-------- 3 files changed, 93 insertions(+), 37 deletions(-) diff --git a/scripts/create_root_files.py b/scripts/create_root_files.py index 1b1c8f2..d960c78 100644 --- a/scripts/create_root_files.py +++ b/scripts/create_root_files.py @@ -3,15 +3,17 @@ run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree: raw////shard-NNN.root -These executables take `[configName] nEvents` 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_events_hits.root, run_sampling writes sampling__events_hits.root, -others may differ again). This script gives each run its own scratch -directory under /.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). +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_events_hits.root, run_sampling writes +sampling__events_hits.root, others may differ again). This script +gives each run its own scratch directory under /.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// — create one first with `dwarf bump-gen`. @@ -105,8 +107,8 @@ def plan_jobs( return jobs -def job_seed(kind: str, gen: str, job: SimJob) -> int: - """Deterministic RNG seed for one sim job, unique per (kind, gen, detector, config, shard). +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 @@ -115,14 +117,31 @@ def job_seed(kind: str, gen: str, job: SimJob) -> int: 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}" + 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, @@ -134,12 +153,9 @@ def run_job( ) workdir.mkdir(parents=True) - cmd = [str(executable)] - if job.config: - cmd.append(job.config) - cmd.append(str(events_per_file)) + cmd = build_cmd(executable, job, events_per_file, energy_gev) - env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job))) + 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: @@ -194,6 +210,7 @@ def run_all( jobs: list[SimJob], executable: Path, events_per_file: int, + energy_gev: float | None, dataset_root: Path, kind: str, gen: str, @@ -208,6 +225,7 @@ def run_all( job, executable, events_per_file, + energy_gev, dataset_root, kind, gen, @@ -238,6 +256,7 @@ def run_make_root( dataset_root: str, jobs: int, execute: bool, + energy_gev: float | None = None, ) -> None: if jobs < 1: raise SystemExit("error: --jobs must be >= 1") @@ -245,6 +264,8 @@ def run_make_root( 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") @@ -257,11 +278,7 @@ def run_make_root( print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===") print(f"executable: {executable}") for job in planned_jobs: - cmd = ( - [str(executable)] - + ([job.config] if job.config else []) - + [str(events_per_file)] - ) + cmd = build_cmd(executable, job, events_per_file, energy_gev) dest = ( dataset_root_path / "raw" @@ -270,7 +287,7 @@ def run_make_root( / job.detector / f"shard-{job.shard_index:03d}.root" ) - seed = job_seed(kind, gen, job) + seed = job_seed(kind, gen, job, energy_gev) print(f" MINICALOSIM_SEED={seed} {' '.join(cmd)} -> {dest}") if not execute: @@ -283,6 +300,7 @@ def run_make_root( planned_jobs, executable, events_per_file, + energy_gev, dataset_root_path, kind, gen, diff --git a/scripts/dwarf.py b/scripts/dwarf.py index e096834..de519f0 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -363,6 +363,16 @@ def make_root( gen: Annotated[ str, typer.Option("--gen", help="Existing gen tag under raw//, e.g. gen1") ], + energy_gev: Annotated[ + float | None, + typer.Option( + "--energy-gev", + help="energy_GeV passed to the executable (default: executable's own " + "default, currently 1.0). Note the dataset detector label is not " + "derived from this — e.g. use '--detector pbwo4_10gev --energy-gev 10' " + "to name the dataset accordingly.", + ), + ] = None, kind: Annotated[ str, typer.Option("--kind", help="steps | hits | ... (default: steps)") ] = "steps", @@ -390,6 +400,7 @@ def make_root( dataset_root=str(dataset_root), jobs=jobs, execute=execute, + energy_gev=energy_gev, ) diff --git a/tests/test_create_root_files.py b/tests/test_create_root_files.py index 4d21f5e..f07453c 100644 --- a/tests/test_create_root_files.py +++ b/tests/test_create_root_files.py @@ -139,24 +139,29 @@ def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path): def test_job_seed_deterministic(): job = SimJob(detector="pbwo4", config=None, shard_index=3) - assert job_seed("steps", "gen1", job) == job_seed("steps", "gen1", job) + assert job_seed("steps", "gen1", job, None) == job_seed("steps", "gen1", job, None) def test_job_seed_varies_by_shard_index(): a = SimJob(detector="pbwo4", config=None, shard_index=0) b = SimJob(detector="pbwo4", config=None, shard_index=1) - assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b) + assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None) def test_job_seed_varies_by_detector(): a = SimJob(detector="pbwo4", config=None, shard_index=0) b = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0) - assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b) + assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None) def test_job_seed_varies_by_gen(): job = SimJob(detector="pbwo4", config=None, shard_index=0) - assert job_seed("steps", "gen1", job) != job_seed("steps", "gen2", job) + assert job_seed("steps", "gen1", job, None) != job_seed("steps", "gen2", job, None) + + +def test_job_seed_varies_by_energy(): + job = SimJob(detector="pbwo4", config=None, shard_index=0) + assert job_seed("steps", "gen1", job, 1.0) != job_seed("steps", "gen1", job, 10.0) def test_run_job_passes_deterministic_seed_env_var(tmp_path): @@ -166,11 +171,11 @@ def test_run_job_passes_deterministic_seed_env_var(tmp_path): tmp_root.mkdir() job = SimJob(detector="pbwo4", config=None, shard_index=5) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert result.dest is not None payload = json.loads(result.dest.read_text()) - assert payload["seed"] == str(job_seed("steps", "gen1", job)) + assert payload["seed"] == str(job_seed("steps", "gen1", job, None)) def test_run_job_moves_output_to_correct_shard_path(tmp_path): @@ -181,7 +186,7 @@ def test_run_job_moves_output_to_correct_shard_path(tmp_path): tmp_root.mkdir() job = SimJob(detector="pbwo4", config=None, shard_index=7) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert result.ok assert result.dest == gen_dir / "pbwo4" / "shard-007.root" @@ -197,7 +202,7 @@ def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path): tmp_root.mkdir() job = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert result.ok assert result.dest is not None @@ -215,13 +220,27 @@ def test_run_job_omits_config_arg_when_none(tmp_path): tmp_root.mkdir() job = SimJob(detector="pbwo4", config=None, shard_index=0) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert result.dest is not None payload = json.loads(result.dest.read_text()) assert payload["argv"] == ["10000"] +def test_run_job_appends_energy_arg_when_given(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py") + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4_10gev", config=None, shard_index=0) + result = run_job(job, fake, 10000, 10.0, tmp_path, "steps", "gen1", tmp_root) + + assert result.dest is not None + payload = json.loads(result.dest.read_text()) + assert payload["argv"] == ["10000", "10.0"] + + def test_run_job_fails_when_executable_errors(tmp_path): fake = _write_fake_executable(tmp_path / "fake_exe.py", exit_code=1) (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) @@ -229,7 +248,7 @@ def test_run_job_fails_when_executable_errors(tmp_path): tmp_root.mkdir() job = SimJob(detector="pbwo4", config=None, shard_index=0) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert not result.ok assert "exited 1" in result.message @@ -242,7 +261,7 @@ def test_run_job_fails_when_no_root_file_produced(tmp_path): tmp_root.mkdir() job = SimJob(detector="pbwo4", config=None, shard_index=0) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert not result.ok assert "found 0" in result.message @@ -255,7 +274,7 @@ def test_run_job_fails_when_multiple_root_files_produced(tmp_path): tmp_root.mkdir() job = SimJob(detector="pbwo4", config=None, shard_index=0) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert not result.ok assert "found 2" in result.message @@ -270,7 +289,7 @@ def test_run_job_refuses_to_overwrite_existing_shard(tmp_path): tmp_root.mkdir() job = SimJob(detector="pbwo4", config=None, shard_index=0) - result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root) assert not result.ok assert "overwrite" in result.message @@ -285,7 +304,15 @@ def test_run_all_caps_concurrency(tmp_path): jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)] results = run_all( - jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root + jobs, + fake, + 10000, + None, + tmp_path, + "steps", + "gen1", + max_workers=2, + tmp_root=tmp_root, ) assert all(r.ok and r.dest is not None for r in results)