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

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:
2026-06-25 16:01:13 +02:00
parent 64c6bd1cef
commit 8475199609
13 changed files with 415 additions and 83 deletions
+4
View File
@@ -26,7 +26,9 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
"material": df["material"].to_numpy(dtype=object),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
"step_length": df["step_length"].to_numpy(dtype=np.float32),
"post_E": df["post_E"].to_numpy(dtype=np.float32),
"delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
"edep": df["edep"].to_numpy(dtype=np.float32),
"post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
@@ -63,6 +65,7 @@ _COND_COLS = [
"material",
"layer_id",
"child_track_ids",
"e_sec",
]
@@ -76,6 +79,7 @@ def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
"material": df["material"].to_numpy(dtype=object),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
}