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 <noreply@anthropic.com>
This commit is contained in:
2026-07-13 13:59:09 +02:00
parent 670f57c309
commit c81076c324
3 changed files with 93 additions and 37 deletions
+41 -14
View File
@@ -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)