Fix column names to match actual parquet schema

material_id → material, n_secondaries derived from child_track_ids.list.len()

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-06-17 14:26:57 +02:00
parent 2e4b4d91d1
commit 78c2a61ecd
2 changed files with 10 additions and 10 deletions
+2 -2
View File
@@ -300,9 +300,9 @@ def predict(
"pre_dir_x": chunk["pre_dir"][:, 0], "pre_dir_x": chunk["pre_dir"][:, 0],
"pre_dir_y": chunk["pre_dir"][:, 1], "pre_dir_y": chunk["pre_dir"][:, 1],
"pre_dir_z": chunk["pre_dir"][:, 2], "pre_dir_z": chunk["pre_dir"][:, 2],
"material_id": chunk["material"], "material": chunk["material"],
"layer_id": chunk["layer_id"], "layer_id": chunk["layer_id"],
"n_secondaries": chunk["n_sec"], "n_sec": chunk["n_sec"],
"step_length": step_length, "step_length": step_length,
"delta_e": delta_e, "delta_e": delta_e,
"edep": edep, "edep": edep,
+8 -8
View File
@@ -23,9 +23,9 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_energy": df["pre_energy"].to_numpy(dtype=np.float32), "pre_energy": df["pre_energy"].to_numpy(dtype=np.float32),
"pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32), "pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32),
"material": df["material_id"].to_numpy(dtype=np.int32), "material": df["material"].to_numpy(dtype=np.int32),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32), "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["n_secondaries"].to_numpy(dtype=np.int32), "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
"step_length": df["step_length"].to_numpy(dtype=np.float32), "step_length": df["step_length"].to_numpy(dtype=np.float32),
"delta_e": (df["pre_energy"] - df["post_energy"]).to_numpy(dtype=np.float32), "delta_e": (df["pre_energy"] - df["post_energy"]).to_numpy(dtype=np.float32),
"edep": df["edep"].to_numpy(dtype=np.float32), "edep": df["edep"].to_numpy(dtype=np.float32),
@@ -53,7 +53,7 @@ _COND_COLS = [
"event_id", "pdg", "event_id", "pdg",
"pre_x", "pre_y", "pre_z", "pre_energy", "pre_x", "pre_y", "pre_z", "pre_energy",
"pre_dir_x", "pre_dir_y", "pre_dir_z", "pre_dir_x", "pre_dir_y", "pre_dir_z",
"material_id", "layer_id", "n_secondaries", "material", "layer_id", "child_track_ids",
] ]
@@ -64,9 +64,9 @@ def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_energy": df["pre_energy"].to_numpy(dtype=np.float32), "pre_energy": df["pre_energy"].to_numpy(dtype=np.float32),
"pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32), "pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32),
"material": df["material_id"].to_numpy(dtype=np.int32), "material": df["material"].to_numpy(dtype=np.int32),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32), "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["n_secondaries"].to_numpy(dtype=np.int32), "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
} }
@@ -91,13 +91,13 @@ def build_index_maps(
def build_index_maps_from_files( def build_index_maps_from_files(
files: list[Path], files: list[Path],
) -> tuple[dict[int, int], dict[int, int]]: ) -> tuple[dict[int, int], dict[int, int]]:
"""Scan only pdg and material_id columns across all files (2-column read).""" """Scan only pdg and material columns across all files (2-column read)."""
pdg_vals: set[int] = set() pdg_vals: set[int] = set()
mat_vals: set[int] = set() mat_vals: set[int] = set()
for path in files: for path in files:
df = pd.read_parquet(path, columns=["pdg", "material_id"]) df = pd.read_parquet(path, columns=["pdg", "material"])
pdg_vals.update(int(v) for v in df["pdg"].unique()) pdg_vals.update(int(v) for v in df["pdg"].unique())
mat_vals.update(int(v) for v in df["material_id"].unique()) mat_vals.update(int(v) for v in df["material"].unique())
return ( return (
{v: i for i, v in enumerate(sorted(pdg_vals))}, {v: i for i, v in enumerate(sorted(pdg_vals))},
{v: i for i, v in enumerate(sorted(mat_vals))}, {v: i for i, v in enumerate(sorted(mat_vals))},