105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow.parquet as pq
|
|
|
|
|
|
def find_parquet_files(path: str | Path) -> list[Path]:
|
|
p = Path(path)
|
|
if p.is_dir():
|
|
files = sorted(p.glob("*.parquet"))
|
|
if not files:
|
|
raise FileNotFoundError(f"no .parquet files found in {p}")
|
|
return files
|
|
return [p]
|
|
|
|
|
|
def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
|
|
return {
|
|
"event_id": df["event_id"].to_numpy(),
|
|
"pdg": df["pdg"].to_numpy(dtype=np.int32),
|
|
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
|
|
"pre_E": df["pre_E"].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=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),
|
|
"delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
|
|
"edep": df["edep"].to_numpy(dtype=np.float32),
|
|
"post_dir": df[["post_dir_x", "post_dir_y", "post_dir_z"]].to_numpy(dtype=np.float32),
|
|
}
|
|
|
|
|
|
def load_steps(path: str | Path) -> dict[str, np.ndarray]:
|
|
return _df_to_dict(pd.read_parquet(path))
|
|
|
|
|
|
def load_event_ids(path: str | Path) -> np.ndarray:
|
|
"""Read only the event_id column — cheap scan for split assignment."""
|
|
return pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
|
|
|
|
|
|
def iter_file_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
|
|
"""Yield one parquet row-group at a time so a large file never fully loads."""
|
|
pf = pq.ParquetFile(path)
|
|
for i in range(pf.num_row_groups):
|
|
yield _df_to_dict(pf.read_row_group(i).to_pandas())
|
|
|
|
|
|
_COND_COLS = [
|
|
"event_id", "pdg",
|
|
"pre_x", "pre_y", "pre_z", "pre_E",
|
|
"pre_dir_x", "pre_dir_y", "pre_dir_z",
|
|
"material", "layer_id", "child_track_ids",
|
|
]
|
|
|
|
|
|
def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
|
|
return {
|
|
"event_id": df["event_id"].to_numpy(),
|
|
"pdg": df["pdg"].to_numpy(dtype=np.int32),
|
|
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
|
|
"pre_E": df["pre_E"].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=object),
|
|
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
|
|
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
|
|
}
|
|
|
|
|
|
def iter_cond_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
|
|
"""Yield conditioning-only row-groups (no post-step columns read from disk)."""
|
|
pf = pq.ParquetFile(path)
|
|
for i in range(pf.num_row_groups):
|
|
yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas())
|
|
|
|
|
|
def build_index_maps(
|
|
data: dict[str, np.ndarray],
|
|
) -> tuple[dict[int, int], dict[str, int]]:
|
|
pdg_vals = sorted(int(v) for v in np.unique(data["pdg"]))
|
|
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)},
|
|
)
|
|
|
|
|
|
def build_index_maps_from_files(
|
|
files: list[Path],
|
|
) -> 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[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(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))},
|
|
)
|