diff --git a/giant/data/loader.py b/giant/data/loader.py index aa7955e..a7be53c 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -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))}, diff --git a/giant/data/transforms.py b/giant/data/transforms.py index bd74970..00fe25b 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -123,7 +123,7 @@ def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> def build_cond_features( data: dict[str, np.ndarray], pdg_map: dict[int, int], - mat_map: dict[int, int], + mat_map: dict[str, int], cond_normalizer: "Normalizer | None" = None, ) -> tuple[np.ndarray, np.ndarray]: """Build conditioning arrays only — no target, no post-step variables.""" @@ -136,7 +136,7 @@ def build_cond_features( ]).astype(np.float32) pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) - mat_idx = np.array([mat_map[int(m)] for m in data["material"]], dtype=np.int64) + mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) cond_cat = np.column_stack([pdg_idx, mat_idx]) if cond_normalizer is not None: @@ -148,7 +148,7 @@ def build_cond_features( def build_features( data: dict[str, np.ndarray], pdg_map: dict[int, int], - mat_map: dict[int, int], + mat_map: dict[str, int], cond_normalizer: Normalizer | None = None, target_normalizer: Normalizer | None = None, fit: bool = False, @@ -175,7 +175,7 @@ def build_features( ]).astype(np.float32) # (N, 9) pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) - mat_idx = np.array([mat_map[int(m)] for m in data["material"]], dtype=np.int64) + mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2) if fit: