Add post_pos as a model target via travel_dir decomposition
step_length already encodes |post_pos - pre_pos| by definition, so a raw post_pos target would duplicate that magnitude and could drift inconsistent with step_length during sampling. Instead add travel_dir, a unit vector (local frame) giving only the direction of pre_pos->post_pos; post_pos is reconstructed at inference as pre_pos + step_length * travel_dir, keeping the two self-consistent. Target grows from 6D to 9D. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -30,6 +30,7 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
|
||||
"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),
|
||||
"post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user