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