lars
4ba419ebe4
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>
46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
import polars as pl
|
|
|
|
from scripts import steps_to_parquet
|
|
|
|
|
|
def _frame() -> pl.DataFrame:
|
|
# event 0: step (1,0) spawns track 2 (first-step pre_E=15) → e_sec=15.
|
|
# event 1: step (1,0) spawns tracks 2 & 3 (20 + 30) → e_sec=50.
|
|
return pl.DataFrame(
|
|
{
|
|
"event_id": [0, 0, 0, 1, 1, 1],
|
|
"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_attributes(_frame())
|
|
e_sec = dict(
|
|
zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
|
|
)
|
|
assert e_sec[(1, 0, 0)] == 15.0 # one child, first-step pre_E 15
|
|
assert e_sec[(1, 0, 1)] == 50.0 # two children, 20 + 30
|
|
|
|
|
|
def test_e_sec_zero_when_no_children():
|
|
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)
|
|
)
|
|
assert childless["e_sec"].item() == 0.0
|
|
|
|
|
|
def test_e_sec_preserves_row_count_and_order():
|
|
df = _frame()
|
|
out = steps_to_parquet._add_secondary_attributes(df)
|
|
assert out.height == df.height
|
|
assert out["pre_E"].to_list() == df["pre_E"].to_list()
|