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:
+2
-2
@@ -300,9 +300,9 @@ def predict(
|
||||
"pre_dir_x": chunk["pre_dir"][:, 0],
|
||||
"pre_dir_y": chunk["pre_dir"][:, 1],
|
||||
"pre_dir_z": chunk["pre_dir"][:, 2],
|
||||
"material_id": chunk["material"],
|
||||
"material": chunk["material"],
|
||||
"layer_id": chunk["layer_id"],
|
||||
"n_secondaries": chunk["n_sec"],
|
||||
"n_sec": chunk["n_sec"],
|
||||
"step_length": step_length,
|
||||
"delta_e": delta_e,
|
||||
"edep": edep,
|
||||
|
||||
@@ -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_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),
|
||||
"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),
|
||||
"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),
|
||||
"delta_e": (df["pre_energy"] - df["post_energy"]).to_numpy(dtype=np.float32),
|
||||
"edep": df["edep"].to_numpy(dtype=np.float32),
|
||||
@@ -53,7 +53,7 @@ _COND_COLS = [
|
||||
"event_id", "pdg",
|
||||
"pre_x", "pre_y", "pre_z", "pre_energy",
|
||||
"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_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),
|
||||
"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),
|
||||
"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(
|
||||
files: list[Path],
|
||||
) -> 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()
|
||||
mat_vals: set[int] = set()
|
||||
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())
|
||||
mat_vals.update(int(v) for v in df["material_id"].unique())
|
||||
mat_vals.update(int(v) for v in df["material"].unique())
|
||||
return (
|
||||
{v: i for i, v in enumerate(sorted(pdg_vals))},
|
||||
{v: i for i, v in enumerate(sorted(mat_vals))},
|
||||
|
||||
Reference in New Issue
Block a user