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:
@@ -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),
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user