Add giant predict command
- iter_cond_chunks: column-projected row-group streaming; post-step variables are never read from disk during inference - build_cond_features: assembles conditioning arrays without any target or post-step fields - inv_local_frame_rotation: Rodrigues R^T (negative angle) to rotate predicted post_dir back from local frame to world frame - giant predict: loads checkpoint, streams input, runs flow matching sampler, inverse-normalises and inverse-rotates outputs, writes predictions incrementally as parquet via PyArrow ParquetWriter - train now saves model_config in checkpoint so predict can reconstruct the architecture without extra CLI flags Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -49,6 +49,34 @@ def iter_file_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
|
||||
yield _df_to_dict(pf.read_row_group(i).to_pandas())
|
||||
|
||||
|
||||
_COND_COLS = [
|
||||
"event_id", "pdg",
|
||||
"pre_x", "pre_y", "pre_z", "pre_energy",
|
||||
"pre_dir_x", "pre_dir_y", "pre_dir_z",
|
||||
"material_id", "layer_id", "n_secondaries",
|
||||
]
|
||||
|
||||
|
||||
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_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_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),
|
||||
}
|
||||
|
||||
|
||||
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[int, int]]:
|
||||
|
||||
Reference in New Issue
Block a user