Merge energy-conservation-poc into phase2-secondary-prediction

Brings the energy-conservation PoC work (dwarf CLI unification, dwarf
status improvements, predict --comment, ODE-step comparison scripts,
predict-parquet-only analysis refactor) onto the Phase 2 branch.

Conflict resolution:
- giant/analysis.py: took the energy-conservation-poc version wholesale.
  That branch deliberately removed the live checkpoint+sampler diagnostics
  path (ModelBundle/load_model_bundle/make_val_loader/collect_samples) in
  favor of reading `giant predict --coord local` parquet output. Phase 2's
  only edits to this file adapted the removed path to the new dataset API,
  so nothing Phase-2-specific is lost; no external code called those funcs.

Fixes for pre-existing breakage surfaced by the merge (both predate it):
- giant/cli.py: predict's `_process` unpacked build_features into 5 values,
  but Phase 2 made it return 8 (added n_sec/sec_cont/sec_pdg_idx). Expanded
  the unpack; `giant predict --coord local` would have crashed otherwise.
- tests/test_steps_to_parquet.py: Phase 2 renamed _add_secondary_energy ->
  _add_secondary_attributes without updating this test. Renamed the calls
  and extended the fixture with the pdg/pre_d{x,y,z} columns the expanded
  function reads; e_sec assertions unchanged.
- analysis/compare_ode_steps_energy_conservation.py: E731 lambda assignment
  (added in the un-linted final PoC commit) rewritten as a def.

ruff, ty, and pytest (179 passed) all green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

tests/test_analysis.py

@@ -84,8 +84,6 @@ def _make_collection(n=200, seed=0, gen_offset=0.0) -> SampleCollection:
         material=rng.choice(["W", "Pb"], size=n),
         real_raw=real,
         gen_raw=gen,
-        real_norm=real,
-        gen_norm=gen,
     )
 
 
@@ -259,8 +257,6 @@ def test_load_predicted_local_round_trips_values(tmp_path):
     np.testing.assert_allclose(
         collection.gen_raw, expected_raw(pred_log_local), atol=1e-4
     )
-    assert collection.real_norm is None
-    assert collection.gen_norm is None
 
 
 def test_load_predicted_local_usable_by_downstream_plots(tmp_path):

tests/test_bump_dataset_version.py

@@ -1,5 +1,4 @@
 import os
-import pytest
 from scripts import bump_dataset_version
 
 plan_bump_gen = bump_dataset_version.plan_bump_gen
@@ -13,7 +12,9 @@ check_holdout_overlap = bump_dataset_version.check_holdout_overlap
 
 
 def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
-    dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
+    dirs, log_line = plan_bump_gen(
+        tmp_path, "steps", "first generation", None, "2026-01-01"
+    )
     assert dirs == [
         tmp_path / "raw" / "steps" / "gen1",
         tmp_path / "processed" / "steps" / "gen1" / "schema1",
@@ -55,14 +56,18 @@ def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
 def test_bump_schema_increments_within_its_gen(tmp_path):
     (tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
     (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
+    dirs, _ = plan_bump_schema(
+        tmp_path, "steps", "gen1", "next schema", None, "2026-01-01"
+    )
     assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]
 
 
 def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
     (tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
     (tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
+    dirs, _ = plan_bump_schema(
+        tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01"
+    )
     assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
 
 
@@ -74,6 +79,50 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
         pass
 
 
+def test_bump_gen_to_specific_tag(tmp_path):
+    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
+    dirs, log_line = plan_bump_gen(
+        tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5"
+    )
+    assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
+    assert "`gen5`" in log_line
+
+
+def test_bump_gen_rejects_invalid_to_tag(tmp_path):
+    try:
+        plan_bump_gen(tmp_path, "steps", "bad tag", None, "2026-01-01", target="v5")
+        assert False, "expected SystemExit"
+    except SystemExit:
+        pass
+
+
+def test_bump_schema_to_specific_tag(tmp_path):
+    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
+    (tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
+    dirs, log_line = plan_bump_schema(
+        tmp_path,
+        "steps",
+        "gen1",
+        "jump to schema5",
+        None,
+        "2026-01-01",
+        target="schema5",
+    )
+    assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema5"]
+    assert "`schema5`" in log_line
+
+
+def test_bump_schema_rejects_invalid_to_tag(tmp_path):
+    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
+    try:
+        plan_bump_schema(
+            tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
+        )
+        assert False, "expected SystemExit"
+    except SystemExit:
+        pass
+
+
 def test_apply_bump_creates_dirs_and_appends_log(tmp_path):
     dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason A", "alice", "2026-01-01")
     apply_bump(tmp_path, dirs, log_line)
@@ -97,6 +146,7 @@ def test_apply_bump_appends_without_clobbering_existing_log(tmp_path):
 # update-manifest
 # ---------------------------------------------------------------------------
 
+
 def _make_parquet(path):
     """Create a zero-byte stand-in for a parquet file."""
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -104,7 +154,15 @@ def _make_parquet(path):
 
 
 def test_update_manifest_bumps_to_specified_schema(tmp_path):
-    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
+    parquet = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema2"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -126,7 +184,15 @@ def test_update_manifest_auto_detects_highest_schema(tmp_path):
     for schema in ("schema1", "schema2", "schema3"):
         d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
         d.mkdir(parents=True)
-    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
+    parquet = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema3"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
     parquet.touch()
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -155,7 +221,15 @@ def test_update_manifest_reports_missing_targets(tmp_path):
 
 
 def test_update_manifest_skips_already_at_target(tmp_path):
-    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
+    parquet = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema2"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -170,7 +244,15 @@ def test_update_manifest_skips_already_at_target(tmp_path):
 
 
 def test_update_manifest_preserves_comments_and_blanks(tmp_path):
-    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
+    parquet = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema2"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -185,8 +267,66 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
     assert lines[2][1] is not None  # the data line was updated
 
 
+def test_update_manifest_bumps_gen(tmp_path):
+    parquet = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen2"
+        / "schema1"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
+    _make_parquet(parquet)
+
+    manifest_dir = tmp_path / "pools" / "pbwo4"
+    manifest_dir.mkdir(parents=True)
+    manifest = manifest_dir / "full.manifest"
+    manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n")
+
+    lines, missing = plan_update_manifest(manifest, None, target_gen="gen2")
+    assert missing == []
+    changed = [(old, new) for old, new in lines if new is not None]
+    assert len(changed) == 1
+    assert "gen2" in changed[0][1]
+    assert "gen1" not in changed[0][1]
+
+
+def test_update_manifest_bumps_gen_and_schema(tmp_path):
+    parquet = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen2"
+        / "schema3"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
+    _make_parquet(parquet)
+
+    manifest_dir = tmp_path / "pools" / "pbwo4"
+    manifest_dir.mkdir(parents=True)
+    manifest = manifest_dir / "full.manifest"
+    manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n")
+
+    lines, missing = plan_update_manifest(manifest, "schema3", target_gen="gen2")
+    assert missing == []
+    changed = [(old, new) for old, new in lines if new is not None]
+    assert len(changed) == 1
+    assert "gen2" in changed[0][1]
+    assert "schema3" in changed[0][1]
+
+
 def test_apply_update_manifest_writes_file(tmp_path):
-    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
+    parquet = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema2"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -206,9 +346,26 @@ def test_apply_update_manifest_writes_file(tmp_path):
 # create-manifest
 # ---------------------------------------------------------------------------
 
+
 def test_create_manifest_writes_relative_paths(tmp_path):
-    pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
-    pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
+    pq1 = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema2"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
+    pq2 = (
+        tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema2"
+        / "pbwo4"
+        / "shard-001.parquet"
+    )
     _make_parquet(pq1)
     _make_parquet(pq2)
 
@@ -217,8 +374,8 @@ def test_create_manifest_writes_relative_paths(tmp_path):
 
     assert missing == []
     assert len(lines) == 2
-    assert all("schema2" in l for l in lines)
-    assert all(not l.startswith("/") for l in lines)
+    assert all("schema2" in line for line in lines)
+    assert all(not line.startswith("/") for line in lines)
     assert resolved == [pq1.resolve(), pq2.resolve()]
 
     apply_create_manifest(output, lines)
@@ -248,6 +405,7 @@ def test_create_manifest_creates_parent_dirs(tmp_path):
 # check_holdout_overlap
 # ---------------------------------------------------------------------------
 
+
 def test_no_overlap_check_when_no_holdout_involved(tmp_path):
     pool_dir = tmp_path / "pools" / "pbwo4"
     pool_dir.mkdir(parents=True)

tests/test_cli_predict.py

@@ -2,7 +2,11 @@ import uuid
 
 import yaml
 
-from giant.cli import _CEPH_PREDICTIONS, _resolve_prediction_output, _write_prediction_ref
+from giant.cli import (
+    _CEPH_PREDICTIONS,
+    _resolve_prediction_output,
+    _write_prediction_ref,
+)
 
 
 # ---------------------------------------------------------------------------
@@ -99,6 +103,25 @@ def test_ref_yaml_contains_expected_fields(tmp_path):
     assert data["dataset"] == str(dataset)
     assert data["checkpoint"] == str(checkpoint.resolve())
     assert "timestamp" in data
+    assert "comment" not in data
+
+
+def test_ref_yaml_includes_comment_when_provided(tmp_path):
+    ckpt_dir = tmp_path / "checkpoints"
+    ckpt_dir.mkdir()
+    checkpoint = ckpt_dir / "best.pt"
+    checkpoint.touch()
+
+    out = tmp_path / "pred.parquet"
+    dataset = tmp_path / "full.manifest"
+    pred_uuid = str(uuid.uuid4())
+
+    ref_path = _write_prediction_ref(
+        checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3"
+    )
+    data = yaml.safe_load(ref_path.read_text())
+
+    assert data["comment"] == "baseline sweep run 3"
 
 
 def test_ref_timestamp_is_iso_format(tmp_path):
@@ -110,7 +133,9 @@ def test_ref_timestamp_is_iso_format(tmp_path):
     checkpoint.touch()
 
     pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
+    ref_path = _write_prediction_ref(
+        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
+    )
     data = yaml.safe_load(ref_path.read_text())
 
     # Must parse without error and be timezone-aware (UTC).
@@ -125,7 +150,9 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
     checkpoint.touch()
 
     pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
+    ref_path = _write_prediction_ref(
+        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
+    )
     data = yaml.safe_load(ref_path.read_text())
 
     assert data["checkpoint"].startswith("/")

tests/test_create_root_files.py

@@ -50,7 +50,10 @@ sys.exit({exit_code})
 
 
 def test_parse_detector_spec_with_config():
-    assert parse_detector_spec("sampling_pb_scint:pb_scint") == ("sampling_pb_scint", "pb_scint")
+    assert parse_detector_spec("sampling_pb_scint:pb_scint") == (
+        "sampling_pb_scint",
+        "pb_scint",
+    )
 
 
 def test_parse_detector_spec_without_config():
@@ -79,13 +82,17 @@ def test_next_shard_index_continues_past_existing(tmp_path):
 
 def test_plan_jobs_rejects_missing_gen(tmp_path):
     with pytest.raises(PlanError):
-        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")
+        plan_jobs(
+            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
+        )
 
 
 def test_plan_jobs_rejects_malformed_gen(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
     with pytest.raises(PlanError):
-        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")
+        plan_jobs(
+            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
+        )
 
 
 def test_plan_jobs_continues_from_existing_shards(tmp_path):
@@ -94,7 +101,9 @@ def test_plan_jobs_continues_from_existing_shards(tmp_path):
     (gen_dir / "pbwo4" / "shard-000.root").touch()
     (gen_dir / "pbwo4" / "shard-001.root").touch()
 
-    jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")
+    jobs = plan_jobs(
+        ["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
+    )
 
     assert [j.shard_index for j in jobs] == [2, 3, 4]
     assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
@@ -133,6 +142,7 @@ def test_run_job_moves_output_to_correct_shard_path(tmp_path):
 
     assert result.ok
     assert result.dest == gen_dir / "pbwo4" / "shard-007.root"
+    assert result.dest is not None
     assert result.dest.is_file()
     assert not any(tmp_root.iterdir())  # workdir cleaned up
 
@@ -147,6 +157,7 @@ def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path):
     result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
 
     assert result.ok
+    assert result.dest is not None
     payload = json.loads(result.dest.read_text())
     assert payload["argv"] == ["pb_scint", "10000"]
     # ran in its own scratch workdir under .sim-tmp, not directly in dataset_root
@@ -163,6 +174,7 @@ def test_run_job_omits_config_arg_when_none(tmp_path):
     job = SimJob(detector="pbwo4", config=None, shard_index=0)
     result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
 
+    assert result.dest is not None
     payload = json.loads(result.dest.read_text())
     assert payload["argv"] == ["10000"]
 
@@ -229,12 +241,15 @@ def test_run_all_caps_concurrency(tmp_path):
     tmp_root.mkdir()
 
     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)
+    results = run_all(
+        jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root
+    )
 
-    assert all(r.ok for r in results)
-    assert {r.dest.name for r in results} == {f"shard-{i:03d}.root" for i in range(6)}
+    assert all(r.ok and r.dest is not None for r in results)
+    dests = [r.dest for r in results if r.dest is not None]
+    assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}
 
-    intervals = [json.loads(r.dest.read_text()) for r in results]
+    intervals = [json.loads(d.read_text()) for d in dests]
     events = sorted(
         [(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
     )
@@ -244,24 +259,3 @@ def test_run_all_caps_concurrency(tmp_path):
         concurrent += delta
         peak = max(peak, concurrent)
     assert peak <= 2
-
-
-def test_build_parser_defaults():
-    args = create_root_files.build_parser().parse_args(
-        [
-            "--executable",
-            "fake",
-            "--detector",
-            "pbwo4",
-            "--num-files",
-            "2",
-            "--events-per-file",
-            "10000",
-            "--gen",
-            "gen1",
-        ]
-    )
-    assert args.jobs == 4
-    assert args.kind == "steps"
-    assert args.dataset_root == "/ceph/lbogner/geant_steps"
-    assert args.execute is False

tests/test_dwarf.py

@@ -0,0 +1,68 @@
+from typer.testing import CliRunner
+
+from scripts.dwarf import app
+
+runner = CliRunner()
+
+
+def test_convert_rejects_jobs_below_one(tmp_path):
+    root_file = tmp_path / "shard.root"
+    root_file.touch()
+    result = runner.invoke(app, ["convert", str(root_file), "--jobs", "0"])
+    assert result.exit_code != 0
+    assert "--jobs must be >= 1" in result.output
+
+
+def test_convert_rejects_output_with_multiple_files(tmp_path):
+    a = tmp_path / "a.root"
+    b = tmp_path / "b.root"
+    a.touch()
+    b.touch()
+    result = runner.invoke(app, ["convert", str(a), str(b), "--output", "out.parquet"])
+    assert result.exit_code != 0
+    assert "--output can only be used with a single input file" in result.output
+
+
+def test_convert_rejects_output_with_parallel_jobs(tmp_path):
+    root_file = tmp_path / "shard.root"
+    root_file.touch()
+    result = runner.invoke(
+        app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
+    )
+    assert result.exit_code != 0
+    assert "--output cannot be combined with --jobs > 1" in result.output
+
+
+def test_convert_default_jobs_is_one():
+    result = runner.invoke(app, ["convert", "--help"])
+    assert result.exit_code == 0
+    assert "default: 1" in result.output
+
+
+def test_bump_gen_requires_reason():
+    result = runner.invoke(app, ["bump-gen"])
+    assert result.exit_code != 0
+    assert "reason" in result.output.lower()
+
+
+def test_create_manifest_requires_exactly_one_of_output_or_pool(tmp_path):
+    f = tmp_path / "a.parquet"
+    f.touch()
+    result = runner.invoke(app, ["create-manifest", str(f)])
+    assert result.exit_code != 0
+    assert "exactly one of --output or --pool is required" in result.output
+
+
+def test_create_manifest_requires_type_with_pool(tmp_path):
+    f = tmp_path / "a.parquet"
+    f.touch()
+    result = runner.invoke(app, ["create-manifest", "--pool", "pbwo4", str(f)])
+    assert result.exit_code != 0
+    assert "--type is required when --pool is given" in result.output
+
+
+def test_status_reports_missing_root(tmp_path):
+    missing = tmp_path / "does-not-exist"
+    result = runner.invoke(app, ["status", "--root", str(missing)])
+    assert result.exit_code != 0
+    assert "is not a directory" in result.output

tests/test_steps_to_parquet.py

@@ -12,13 +12,17 @@ def _frame() -> pl.DataFrame:
             "track_id": [1, 1, 2, 1, 2, 3],
             "step_no": [0, 1, 0, 0, 0, 0],
             "pre_E": [100.0, 80.0, 15.0, 200.0, 20.0, 30.0],
+            "pdg": [11, 11, 22, 11, 22, 22],
+            "pre_dx": [0.0, 0.0, 1.0, 0.0, 1.0, 0.0],
+            "pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+            "pre_dz": [1.0, 1.0, 0.0, 1.0, 0.0, 0.0],
             "child_track_ids": [[2], [], [], [2, 3], [], []],
         }
     )
 
 
 def test_e_sec_sums_child_first_step_energy():
-    out = steps_to_parquet._add_secondary_energy(_frame())
+    out = steps_to_parquet._add_secondary_attributes(_frame())
     e_sec = dict(
         zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
     )
@@ -27,7 +31,7 @@ def test_e_sec_sums_child_first_step_energy():
 
 
 def test_e_sec_zero_when_no_children():
-    out = steps_to_parquet._add_secondary_energy(_frame())
+    out = steps_to_parquet._add_secondary_attributes(_frame())
     childless = out.filter(
         (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
     )
@@ -36,6 +40,6 @@ def test_e_sec_zero_when_no_children():
 
 def test_e_sec_preserves_row_count_and_order():
     df = _frame()
-    out = steps_to_parquet._add_secondary_energy(df)
+    out = steps_to_parquet._add_secondary_attributes(df)
     assert out.height == df.height
     assert out["pre_E"].to_list() == df["pre_E"].to_list()

tests/test_steps_to_parquet_parallel.py

@@ -1,4 +1,5 @@
 import json
+import sys
 from pathlib import Path
 
 from scripts import steps_to_parquet_parallel
@@ -40,7 +41,7 @@ def test_runs_one_job_per_file_and_reports_success(tmp_path):
     fake = _write_fake_executable(tmp_path, marker_dir)
 
     files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(3)]
-    results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)
+    results = run_parallel(files, jobs=4, cmd_prefix=[sys.executable, str(fake)])
 
     assert {r[0] for r in results} == set(files)
     assert all(code == 0 for _, code, _, _ in results)
@@ -53,7 +54,7 @@ def test_failures_are_reported_with_nonzero_exit_code(tmp_path):
     fake = _write_fake_executable(tmp_path, marker_dir)
 
     files = [str(tmp_path / "shard-000.root"), str(tmp_path / "shard-fail.root")]
-    results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)
+    results = run_parallel(files, jobs=4, cmd_prefix=[sys.executable, str(fake)])
 
     codes = {Path(f).stem: code for f, code, _, _ in results}
     assert codes["shard-000"] == 0
@@ -66,7 +67,7 @@ def test_jobs_caps_concurrency(tmp_path):
     fake = _write_fake_executable(tmp_path, marker_dir)
 
     files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(6)]
-    run_parallel(files, jobs=2, steps_to_parquet_path=fake)
+    run_parallel(files, jobs=2, cmd_prefix=[sys.executable, str(fake)])
 
     intervals = []
     for f in files:
@@ -83,11 +84,6 @@ def test_jobs_caps_concurrency(tmp_path):
     assert peak <= 2
 
 
-def test_default_jobs_is_four():
-    args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
-    assert args.jobs == 4
-
-
 def test_output_for_is_passed_through_as_output_flag(tmp_path):
     marker_dir = tmp_path / "markers"
     marker_dir.mkdir()
@@ -108,7 +104,10 @@ sys.exit(0)
     root_file = str(tmp_path / "shard-000.root")
     dest = tmp_path / "processed" / "shard-000.parquet"
     run_parallel(
-        [root_file], jobs=1, steps_to_parquet_path=fake, output_for={root_file: dest}
+        [root_file],
+        jobs=1,
+        cmd_prefix=[sys.executable, str(fake)],
+        output_for={root_file: dest},
     )
     assert (marker_dir / "shard-000.txt").read_text() == str(dest)
 
@@ -125,13 +124,31 @@ def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
 def test_resolve_destination_uses_latest_schema(tmp_path):
     root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
     dest = resolve_destination(root_file, tmp_path, schema_override=None)
-    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
+    assert (
+        dest
+        == tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema3"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
 
 
 def test_resolve_destination_schema_override_wins(tmp_path):
     root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
     dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
-    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"
+    assert (
+        dest
+        == tmp_path
+        / "processed"
+        / "steps"
+        / "gen1"
+        / "schema9"
+        / "pbwo4"
+        / "shard-000.parquet"
+    )
 
 
 def test_resolve_destination_errors_without_any_schema(tmp_path):
@@ -171,9 +188,3 @@ def test_resolve_destination_errors_on_wrong_shape(tmp_path):
 
 def test_latest_schema_tag_returns_none_when_missing(tmp_path):
     assert latest_schema_tag(tmp_path / "does" / "not" / "exist") is None
-
-
-def test_dataset_root_and_schema_flags_default(tmp_path):
-    args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
-    assert args.dataset_root == "/ceph/lbogner/geant_steps"
-    assert args.schema is None

tests/test_transforms.py

@@ -1,7 +1,9 @@
 import numpy as np
+import pytest
 from giant.data.transforms import (
     energy_simplex_decode,
     energy_simplex_encode,
+    inv_local_frame_rotation,
     inv_log_transform,
     local_frame_rotation,
     log_transform,
@@ -53,6 +55,35 @@ def test_local_frame_rotation_preserves_norm():
     np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)
 
 
+def test_local_frame_rotation_rejects_near_zero_pre_dir():
+    """A degenerate (near-zero-norm) pre_dir has no well-defined frame — must
+    raise instead of silently falling back to an arbitrary rotation axis."""
+    pre_dir = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
+    post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
+    with pytest.raises(ValueError, match="near-zero norm"):
+        local_frame_rotation(pre_dir, post_dir)
+    with pytest.raises(ValueError, match="near-zero norm"):
+        inv_local_frame_rotation(pre_dir, post_dir)
+
+
+def test_local_frame_rotation_normalizes_non_unit_pre_dir():
+    """A pre_dir with float32-drift norm (not exactly 1) must still produce the
+    same result as its exactly-normalized counterpart, not a skewed frame."""
+    rng = np.random.default_rng(9)
+    N = 50
+    pre_dir_unit = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir_unit /= np.linalg.norm(pre_dir_unit, axis=1, keepdims=True)
+    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
+    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
+
+    pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(
+        np.float32
+    )
+    expected = local_frame_rotation(pre_dir_unit, post_dir)
+    result = local_frame_rotation(pre_dir_scaled, post_dir)
+    np.testing.assert_allclose(result, expected, atol=1e-4)
+
+
 def test_travel_direction_is_unit_norm():
     rng = np.random.default_rng(5)
     N = 50
@@ -128,7 +159,7 @@ def test_energy_simplex_conservation():
 def test_energy_simplex_roundtrip():
     """Encode → decode recovers energies whose lost part already sums to delta_e."""
     rng = np.random.default_rng(12)
-    N = 500
+    N = 500000
     pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
     post_E = (pre_E * rng.uniform(0.0, 1.0, N)).astype(np.float32)
     delta_e = pre_E - post_E
@@ -170,5 +201,7 @@ def test_normalizer_serialization():
     X = rng.standard_normal((50, 6)).astype(np.float32)
     norm = Normalizer().fit(X)
     norm2 = Normalizer.from_dict(norm.to_dict())
+    assert norm2.mean is not None and norm.mean is not None
+    assert norm2.std is not None and norm.std is not None
     np.testing.assert_allclose(norm2.mean, norm.mean)
     np.testing.assert_allclose(norm2.std, norm.std)