lars
05d5dee606
The merged proc_idx/proc_map plumbing wasn't run through ruff format before merging; reflow only, no logic changes.
228 lines
8.4 KiB
Python
228 lines
8.4 KiB
Python
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow.parquet as pq
|
|
|
|
# A manifest is a plain text file listing one parquet path per line, used to
|
|
# name a curated subset of files (e.g. a train/holdout pool) without copying
|
|
# or symlinking the underlying parquet files. Lines are resolved relative to
|
|
# the manifest's own directory, so the manifest stays valid if the whole
|
|
# dataset tree is moved or copied elsewhere intact.
|
|
MANIFEST_SUFFIX = ".manifest"
|
|
|
|
|
|
def _read_manifest(path: Path) -> list[Path]:
|
|
files = []
|
|
for line in path.read_text().splitlines():
|
|
line = line.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
resolved = (path.parent / line).resolve()
|
|
if not resolved.is_file():
|
|
raise FileNotFoundError(f"{path} lists missing file: {resolved}")
|
|
files.append(resolved)
|
|
if not files:
|
|
raise FileNotFoundError(f"manifest {path} lists no files")
|
|
return files
|
|
|
|
|
|
def find_parquet_files(path: str | Path) -> list[Path]:
|
|
p = Path(path)
|
|
if p.suffix == MANIFEST_SUFFIX:
|
|
return _read_manifest(p)
|
|
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 _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
|
|
"""Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
|
|
out = np.full((len(series), K), fill, dtype=np.float32)
|
|
for i, lst in enumerate(series):
|
|
if lst is not None and len(lst) > 0:
|
|
n = min(len(lst), K)
|
|
out[i, :n] = lst[:n]
|
|
return out
|
|
|
|
|
|
def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
|
|
"""Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
|
|
out = np.full((len(series), K), fill, dtype=np.int64)
|
|
for i, lst in enumerate(series):
|
|
if lst is not None and len(lst) > 0:
|
|
n = min(len(lst), K)
|
|
out[i, :n] = lst[:n]
|
|
return out
|
|
|
|
|
|
def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
|
|
"""Pad three list-valued direction columns → (N, K, 3) float32.
|
|
|
|
Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
|
|
"""
|
|
N = len(dx)
|
|
out = np.zeros((N, K, 3), dtype=np.float32)
|
|
out[:, :, 2] = 1.0
|
|
for i in range(N):
|
|
lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
|
|
if lx is not None and len(lx) > 0:
|
|
n = min(len(lx), K)
|
|
out[i, :n, 0] = lx[:n]
|
|
out[i, :n, 1] = ly[:n]
|
|
out[i, :n, 2] = lz[:n]
|
|
return out
|
|
|
|
|
|
def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
|
|
from giant.constants import K_MAX
|
|
|
|
has_sec_lists = "sec_E_list" in df.columns
|
|
|
|
d: dict[str, np.ndarray] = {
|
|
"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_dx", "pre_dy", "pre_dz"]].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),
|
|
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
|
|
# The physics process that ended the step (e.g. "compt", "phot",
|
|
# "eBrem") — a post-step outcome, so it's a router/classifier
|
|
# supervision label only, never conditioning (see build_process_map*
|
|
# / ProcessRouter). Guarded like has_sec_lists: older parquet
|
|
# conversions predating this column still load fine.
|
|
"process": (
|
|
df["process"].to_numpy(dtype=object)
|
|
if "process" in df.columns
|
|
else np.full(len(df), "", dtype=object)
|
|
),
|
|
"step_length": df["step_length"].to_numpy(dtype=np.float32),
|
|
"post_E": df["post_E"].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_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
|
|
"post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
|
|
}
|
|
|
|
if has_sec_lists:
|
|
d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
|
|
d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
|
|
d["sec_dir_list"] = _pad_dir_col(
|
|
df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
|
|
)
|
|
|
|
return d
|
|
|
|
|
|
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_dx",
|
|
"pre_dy",
|
|
"pre_dz",
|
|
"material",
|
|
"layer_id",
|
|
"child_track_ids",
|
|
"e_sec",
|
|
]
|
|
|
|
|
|
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_dx", "pre_dy", "pre_dz"]].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),
|
|
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
|
|
}
|
|
|
|
|
|
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))},
|
|
)
|
|
|
|
|
|
def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
|
|
"""Scan the `process` column and build a frequency-capped process->index map.
|
|
|
|
Physics processes have a long tail (rare nuclear captures, decays, ...)
|
|
while `ProcessRouter` needs a fixed number of expert slots, so only the
|
|
`n_experts - 1` most frequent processes get their own index; every rarer
|
|
process is bucketed into a shared "other" index (`n_experts - 1`). This
|
|
mirrors how `build_features` clamps the n_sec label to K_MAX for the
|
|
fixed-width n_sec_head classifier.
|
|
"""
|
|
counts: dict[str, int] = {}
|
|
for path in files:
|
|
df = pd.read_parquet(path, columns=["process"])
|
|
for name, count in df["process"].value_counts().items():
|
|
name = str(name)
|
|
counts[name] = counts.get(name, 0) + int(count)
|
|
ranked = sorted(counts, key=lambda name: counts[name], reverse=True)
|
|
keep = ranked[: max(n_experts - 1, 0)]
|
|
proc_map = {name: i for i, name in enumerate(keep)}
|
|
other_idx = n_experts - 1
|
|
for name in ranked[len(keep) :]:
|
|
proc_map[name] = other_idx
|
|
return proc_map
|