lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Encode edep/secondary/post energy as a conservation-constrained simplex

Browse Source 
Replaces the independent log_delta_e/log_edep targets with 2 additive-log-ratio
coordinates over the deposit/secondary/post-energy simplex (fractions of pre_E
summing to 1), so edep + e_sec + post_E == pre_E holds by construction after
decoding (softmax) rather than being learned approximately. Requires e_sec
(secondary energy) as a new conditioning input and a steps_to_parquet.py pass
to derive it from child track first-step energies.
This commit is contained in:
lars
2026-06-25 16:01:13 +02:00
parent 64c6bd1cef
commit 8475199609
13 changed files with 415 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_analysis.py
+50 -20
View File
@@ -41,7 +41,12 @@ from giant.constants import (
PREDICT_SCHEMA_VERSION,
PREDICT_SCHEMA_VERSION_KEY,
)
from giant.data.transforms import inv_log_transform, log_transform, reconstruct_post_pos
from giant.data.transforms import (
energy_simplex_decode,
inv_log_transform,
log_transform,
reconstruct_post_pos,
)
def _unit_vectors(rng, n):
@@ -186,13 +191,17 @@ def test_plot_kl_bars_caps_groups_by_pdg():
def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
"""Mimic `giant predict --coord local`'s output schema for the loader tests."""
"""Mimic `giant predict --coord local`'s output schema for the loader tests.
Column 0 is a log-scaled step_length; columns 1–2 are the deposit/secondary
ALR energy logits (unconstrained reals, decoded against pre_E); columns 3–8
are direction components.
"""
rng = rng or np.random.default_rng(0)
true_log_local = rng.standard_normal((n, 9)).astype(np.float32)
true_log_local[:, :3] = log_transform(
rng.uniform(0.1, 5.0, (n, 3)).astype(np.float32)
)
true_log_local[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
pred_log_local = true_log_local + rng.normal(0, 0.01, (n, 9)).astype(np.float32)
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
table = pa.table(
{
@@ -201,7 +210,7 @@ def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
"pre_x": rng.standard_normal(n).astype(np.float32),
"pre_y": rng.standard_normal(n).astype(np.float32),
"pre_z": rng.standard_normal(n).astype(np.float32),
"pre_E": rng.uniform(1.0, 100.0, n).astype(np.float32),
"pre_E": pre_E,
"pre_dx": rng.standard_normal(n).astype(np.float32),
"pre_dy": rng.standard_normal(n).astype(np.float32),
"pre_dz": rng.standard_normal(n).astype(np.float32),
@@ -221,12 +230,12 @@ def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
if metadata is not None:
table = table.replace_schema_metadata(metadata)
pq.write_table(table, path)
return true_log_local, pred_log_local
return true_log_local, pred_log_local, pre_E
def test_load_predicted_local_round_trips_values(tmp_path):
path = tmp_path / "predicted_local.parquet"
true_log_local, pred_log_local = _write_predicted_local_parquet(
true_log_local, pred_log_local, pre_E = _write_predicted_local_parquet(
path,
metadata={
PREDICT_COORD_METADATA_KEY: "local",
@@ -236,12 +245,20 @@ def test_load_predicted_local_round_trips_values(tmp_path):
collection = load_predicted_local(path)
expected_real = true_log_local.copy()
expected_real[:, :3] = np.exp(expected_real[:, :3]) - 1e-8
expected_gen = pred_log_local.copy()
expected_gen[:, :3] = np.exp(expected_gen[:, :3]) - 1e-8
np.testing.assert_allclose(collection.real_raw, expected_real, atol=1e-4)
np.testing.assert_allclose(collection.gen_raw, expected_gen, atol=1e-4)
def expected_raw(log_local):
raw = log_local.copy()
raw[:, 0] = np.exp(log_local[:, 0]) - 1e-8
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(log_local[:, 1:3], pre_E)
raw[:, 1] = delta_e
raw[:, 2] = edep
return raw
np.testing.assert_allclose(
collection.real_raw, expected_raw(true_log_local), atol=1e-4
)
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
@@ -321,13 +338,21 @@ def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
assert list(expected["group"]) == list(actual["group"])
assert list(expected["n"]) == list(actual["n"])
for col in ["real_mean", "gen_mean", "real_std", "gen_std", "kl_real_gen"]:
for col in ["real_mean", "gen_mean", "real_std", "gen_std"]:
np.testing.assert_allclose(
expected[col].to_numpy(),
actual[col].to_numpy(),
atol=1e-4,
rtol=1e-4,
)
# KL uses np.histogram (numpy path) vs polars Series.hist (lazy path); the two
# backends bin the boundary (min/max) sample differently, so allow a small
# absolute discrepancy rather than requiring bit-identical estimates.
np.testing.assert_allclose(
expected["kl_real_gen"].to_numpy(),
actual["kl_real_gen"].to_numpy(),
atol=2e-2,
)
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
@@ -387,7 +412,8 @@ def _make_event_level_arrays(rng):
def _local_block():
block = rng.standard_normal((n, 9)).astype(np.float32)
block[:, :3] = log_transform(rng.uniform(0.1, 5.0, (n, 3)).astype(np.float32))
block[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
# cols 1–2 stay as random ALR energy logits (decoded against pre_E)
block[:, 3:6] = _unit_vectors(rng, n)
block[:, 6:9] = _unit_vectors(rng, n)
return block
@@ -451,7 +477,9 @@ def _expected_event_table(
def agg(log_local):
step_length = inv_log_transform(log_local[mask, 0])
edep = inv_log_transform(log_local[mask, 2])
edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(
log_local[mask, 1:3], pre_E[mask]
)
travel_dir_local = log_local[mask, 6:9]
post_pos = reconstruct_post_pos(
pre_pos[mask], pre_dir[mask], step_length, travel_dir_local
@@ -574,10 +602,12 @@ def test_event_level_plots_run_without_error(tmp_path):
def test_pdg_contribution_table_pl_matches_manual_sums(tmp_path):
path = tmp_path / "event_level.parquet"
_, _, _, _, true_log_local, pred_log_local, pdg = _write_event_level_parquet(path)
_, _, _, pre_E, true_log_local, pred_log_local, pdg = _write_event_level_parquet(
path
)
real_edep = inv_log_transform(true_log_local[:, 2])
gen_edep = inv_log_transform(pred_log_local[:, 2])
real_edep = energy_simplex_decode(true_log_local[:, 1:3], pre_E)[0]
gen_edep = energy_simplex_decode(pred_log_local[:, 1:3], pre_E)[0]
real_length = inv_log_transform(true_log_local[:, 0])
gen_length = inv_log_transform(pred_log_local[:, 0])
tests/test_flow.py
+4 -3
View File
@@ -1,4 +1,5 @@
import torch
from giant.constants import COND_DIM
from giant.model.network import DenoisingMLP
from giant.model.schedule import CosineSchedule, flow_matching_loss
from giant.sample import sample_flow, sample_ddim
@@ -10,7 +11,7 @@ def _small_model():
def _batch(B=8):
x1 = torch.randn(B, 9)
cond_cont = torch.randn(B, 9)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
return x1, cond_cont, cond_cat
@@ -36,7 +37,7 @@ def test_flow_matching_loss_has_grad():
def test_sample_flow_shape():
B = 6
cond_cont = torch.randn(B, 9)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
out = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
assert out.shape == (B, 9)
@@ -52,7 +53,7 @@ def test_ddpm_loss_nonneg():
def test_sample_ddim_shape():
B = 4
schedule = CosineSchedule(T=50)
cond_cont = torch.randn(B, 9)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
out = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
assert out.shape == (B, 9)
tests/test_network.py
+3 -2
View File
@@ -1,4 +1,5 @@
import torch
from giant.constants import COND_DIM
from giant.model.network import DenoisingMLP, SinusoidalEmbedding
@@ -19,7 +20,7 @@ def test_denoising_mlp_output_shape():
model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
x_t = torch.randn(B, 9)
t = torch.rand(B)
cond_cont = torch.randn(B, 9)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[
torch.randint(0, 5, (B,)),
@@ -36,7 +37,7 @@ def test_denoising_mlp_gradients_flow():
model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
x_t = torch.randn(B, 9)
t = torch.rand(B)
cond_cont = torch.randn(B, 9)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
loss = model(x_t, t, cond_cont, cond_cat).sum()
loss.backward()
tests/test_steps_to_parquet.py
+50
View File
@@ -0,0 +1,50 @@
import importlib.util
from pathlib import Path
import polars as pl
# scripts/ is not an installed package — load the module straight from its path.
_SPEC = importlib.util.spec_from_file_location(
"steps_to_parquet",
Path(__file__).resolve().parents[1] / "scripts" / "steps_to_parquet.py",
)
steps_to_parquet = importlib.util.module_from_spec(_SPEC)
_SPEC.loader.exec_module(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],
"child_track_ids": [[2], [], [], [2, 3], [], []],
}
)
def test_e_sec_sums_child_first_step_energy():
out = steps_to_parquet._add_secondary_energy(_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_energy(_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_energy(df)
assert out.height == df.height
assert out["pre_E"].to_list() == df["pre_E"].to_list()
tests/test_transforms.py
+47
View File
@@ -1,5 +1,7 @@
import numpy as np
from giant.data.transforms import (
energy_simplex_decode,
energy_simplex_encode,
inv_log_transform,
local_frame_rotation,
log_transform,
@@ -111,6 +113,51 @@ def test_reconstruct_post_pos_general_roundtrip():
np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
def test_energy_simplex_conservation():
"""Decoding any ALR coords yields energies that sum to pre_E exactly."""
rng = np.random.default_rng(11)
N = 500
z = rng.standard_normal((N, 2)).astype(np.float32) * 3.0
pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
edep, e_sec, post_E, delta_e = energy_simplex_decode(z, pre_E)
np.testing.assert_allclose(edep + e_sec + post_E, pre_E, rtol=1e-5, atol=1e-4)
np.testing.assert_allclose(delta_e, edep + e_sec, rtol=1e-5, atol=1e-4)
assert np.all(edep >= 0) and np.all(e_sec >= 0) and np.all(post_E >= 0)
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
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
g = rng.uniform(0.0, 1.0, N).astype(np.float32)
edep = (g * delta_e).astype(np.float32)
e_sec = ((1.0 - g) * delta_e).astype(np.float32)
z = energy_simplex_encode(edep, e_sec, post_E, pre_E)
edep_r, e_sec_r, post_E_r, _ = energy_simplex_decode(z, pre_E)
# Tolerance reflects the tiny simplex floor (~1e-5 of pre_E).
np.testing.assert_allclose(edep_r, edep, atol=5e-3)
np.testing.assert_allclose(e_sec_r, e_sec, atol=5e-3)
np.testing.assert_allclose(post_E_r, post_E, atol=5e-3)
def test_energy_simplex_handles_boundary_zeros():
"""e_sec=0 (no secondaries) and post_E=0 (track end) stay finite and decode near 0."""
pre_E = np.array([10.0, 50.0, 100.0], dtype=np.float32)
edep = np.array([4.0, 50.0, 0.0], dtype=np.float32)
e_sec = np.array([0.0, 0.0, 0.0], dtype=np.float32) # no secondaries
post_E = np.array([6.0, 0.0, 100.0], dtype=np.float32) # row 1: track ends
z = energy_simplex_encode(edep, e_sec, post_E, pre_E)
assert np.all(np.isfinite(z))
_, e_sec_r, post_E_r, _ = energy_simplex_decode(z, pre_E)
np.testing.assert_allclose(e_sec_r, 0.0, atol=1e-2)
assert post_E_r[1] < 1e-2 # the absorbed track decodes to ~0 post energy
def test_normalizer_roundtrip():
rng = np.random.default_rng(3)
X = rng.standard_normal((200, 9)).astype(np.float32)