Handle material column as string type
material is a string literal (e.g. "G4_PbWO4"), not an integer. Store as object array and key mat_map on str throughout loader and transforms. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -23,7 +23,7 @@ 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"].to_numpy(dtype=np.int32),
|
||||
"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),
|
||||
"step_length": df["step_length"].to_numpy(dtype=np.float32),
|
||||
@@ -64,7 +64,7 @@ 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"].to_numpy(dtype=np.int32),
|
||||
"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),
|
||||
}
|
||||
@@ -79,9 +79,9 @@ def iter_cond_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
|
||||
|
||||
def build_index_maps(
|
||||
data: dict[str, np.ndarray],
|
||||
) -> tuple[dict[int, int], dict[int, int]]:
|
||||
) -> tuple[dict[int, int], dict[str, int]]:
|
||||
pdg_vals = sorted(int(v) for v in np.unique(data["pdg"]))
|
||||
mat_vals = sorted(int(v) for v in np.unique(data["material"]))
|
||||
mat_vals = sorted(str(v) for v in np.unique(data["material"]))
|
||||
return (
|
||||
{v: i for i, v in enumerate(pdg_vals)},
|
||||
{v: i for i, v in enumerate(mat_vals)},
|
||||
@@ -90,14 +90,14 @@ def build_index_maps(
|
||||
|
||||
def build_index_maps_from_files(
|
||||
files: list[Path],
|
||||
) -> tuple[dict[int, int], dict[int, int]]:
|
||||
) -> tuple[dict[int, int], dict[str, int]]:
|
||||
"""Scan only pdg and material columns across all files (2-column read)."""
|
||||
pdg_vals: set[int] = set()
|
||||
mat_vals: set[int] = set()
|
||||
mat_vals: set[str] = set()
|
||||
for path in files:
|
||||
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"].unique())
|
||||
mat_vals.update(str(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