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:
@@ -3,15 +3,17 @@ 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` 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).
|
||||
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`.
|
||||
@@ -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,
|
||||
|
||||
@@ -363,6 +363,16 @@ def make_root(
|
||||
gen: Annotated[
|
||||
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, 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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user