From a14a4f973ab190dcb86e18f4bc2d4dd8f6b4898c Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 8 Jul 2026 14:44:53 +0200 Subject: [PATCH] Apply ruff format across the codebase Whitespace-only reflow (line wrapping, blank lines between defs); no logic changes. Co-Authored-By: Claude Opus 4.8 --- .../compare_ode_steps_energy_conservation.py | 41 +++- analysis/compare_ode_steps_kl.py | 4 +- analysis/export_energy_conservation_poc.py | 8 +- analysis/export_rollout_observables.py | 4 +- giant/analysis.py | 68 ++++-- giant/cli.py | 28 ++- giant/constants.py | 2 +- giant/data/dataset.py | 29 ++- giant/data/loader.py | 4 +- giant/data/transforms.py | 4 +- giant/model/network.py | 8 +- giant/pipeline.py | 2 +- giant/rollout.py | 227 ++++++++++++++---- giant/sample.py | 4 +- tests/test_phase2.py | 17 +- tests/test_rollout.py | 25 +- 16 files changed, 351 insertions(+), 124 deletions(-) diff --git a/analysis/compare_ode_steps_energy_conservation.py b/analysis/compare_ode_steps_energy_conservation.py index 6ab6633..c78fe86 100644 --- a/analysis/compare_ode_steps_energy_conservation.py +++ b/analysis/compare_ode_steps_energy_conservation.py @@ -38,9 +38,7 @@ def per_event(file: str) -> dict: E0 = float(primary_E[0]) real = pe["real_total_edep"].to_numpy() gen = pe["gen_total_edep"].to_numpy() - n_steps = ( - pl.scan_parquet(file).select(pl.len()).collect(engine="streaming").item() - ) + n_steps = pl.scan_parquet(file).select(pl.len()).collect(engine="streaming").item() return { "E0": E0, "n_events": pe.height, @@ -78,6 +76,8 @@ fmt_row("gen mean/E0", lambda r: f"{r['gen'].mean() / r['E0']:.4f}") fmt_row("gen max/E0", lambda r: f"{r['gen'].max() / r['E0']:.4f}") fmt_row("gen p99/E0", lambda r: f"{np.quantile(r['gen'], 0.99) / r['E0']:.4f}") fmt_row("frac events gen>E0", lambda r: f"{np.mean(r['gen'] > r['E0']):.4f}") + + def disp_ratio(r): return (r["gen"].std() / r["gen"].mean()) / (r["real"].std() / r["real"].mean()) @@ -92,21 +92,31 @@ E0, real_tot, gen_tot = r["E0"], r["real"], r["gen"] fig, ax = plt.subplots(figsize=(6, 4)) edges = _hist_edges(real_tot, gen_tot, bins=50).tolist() ax.hist( - real_tot, bins=edges, density=True, histtype="step", + real_tot, + bins=edges, + density=True, + histtype="step", label=f"real (σ/μ={real_tot.std() / real_tot.mean():.3f})", ) ax.hist( - gen_tot, bins=edges, density=True, histtype="step", + gen_tot, + bins=edges, + density=True, + histtype="step", label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, " f"{np.mean(gen_tot > E0):.1%} > E0)", ) -ax.axvline(E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV") +ax.axvline( + E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV" +) ax.set_yscale("log") ax.set_xlabel("total deposited energy per event [MeV]") ax.set_title("20 ODE steps") ax.legend(fontsize=8) fig.tight_layout() -fig.savefig(OUT / f"{PREFIX20}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight") +fig.savefig( + OUT / f"{PREFIX20}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight" +) fig, ax = plt.subplots(figsize=(6, 4)) ratio_real = real_tot / E0 @@ -129,12 +139,19 @@ all_arrays = [results["10-step (baseline)"]["real"]] + [ ] edges = _hist_edges(*all_arrays, bins=60).tolist() ax.hist( - results["20-step"]["real"], bins=edges, density=True, histtype="step", - color="k", label="real", + results["20-step"]["real"], + bins=edges, + density=True, + histtype="step", + color="k", + label="real", ) for name, r2 in results.items(): ax.hist( - r2["gen"], bins=edges, density=True, histtype="step", + r2["gen"], + bins=edges, + density=True, + histtype="step", label=f"gen {name} ({np.mean(r2['gen'] > r2['E0']):.1%} > E0)", ) ax.axvline(E0, color="gray", linestyle="--", linewidth=1, label=f"E0={E0:.0f} MeV") @@ -143,6 +160,8 @@ ax.set_xlabel("total deposited energy per event [MeV]") ax.set_title("Generated event energy: 10 vs 20 ODE steps") ax.legend(fontsize=8) fig.tight_layout() -fig.savefig(OUT / f"{PREFIX20}-compare-event-total-energy.png", dpi=150, bbox_inches="tight") +fig.savefig( + OUT / f"{PREFIX20}-compare-event-total-energy.png", dpi=150, bbox_inches="tight" +) print("DONE") diff --git a/analysis/compare_ode_steps_kl.py b/analysis/compare_ode_steps_kl.py index e70b617..743e857 100644 --- a/analysis/compare_ode_steps_kl.py +++ b/analysis/compare_ode_steps_kl.py @@ -66,6 +66,4 @@ print( f"{'SUM':<14}{tot['10-step']:>14.5f}{tot['20-step']:>14.5f}" f"{tot['20-step'] / tot['10-step']:>14.2f}" ) -print( - f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}" -) +print(f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}") diff --git a/analysis/export_energy_conservation_poc.py b/analysis/export_energy_conservation_poc.py index 0e72708..d2119ea 100644 --- a/analysis/export_energy_conservation_poc.py +++ b/analysis/export_energy_conservation_poc.py @@ -77,12 +77,16 @@ ax.hist( label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, " f"{np.mean(gen_tot > E0):.1%} > E0)", ) -ax.axvline(E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV") +ax.axvline( + E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV" +) ax.set_yscale("log") ax.set_xlabel("total deposited energy per event [MeV]") ax.legend(fontsize=8) fig.tight_layout() -fig.savefig(OUT / f"{PREFIX}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight") +fig.savefig( + OUT / f"{PREFIX}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight" +) print("=== plot: total edep / incident energy ratio ===") fig, ax = plt.subplots(figsize=(6, 4)) diff --git a/analysis/export_rollout_observables.py b/analysis/export_rollout_observables.py index f54e143..ddad101 100644 --- a/analysis/export_rollout_observables.py +++ b/analysis/export_rollout_observables.py @@ -13,7 +13,9 @@ from pathlib import Path import giant.analysis as a rollout_file = sys.argv[1] if len(sys.argv) > 1 else "rollout.parquet" -reference_file = sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] not in ("", "-") else None +reference_file = ( + sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] not in ("", "-") else None +) OUT = Path(sys.argv[3]) if len(sys.argv) > 3 else Path(".") print(f"=== computing rollout observables: {rollout_file} ===") diff --git a/giant/analysis.py b/giant/analysis.py index c8c110b..45ad39c 100644 --- a/giant/analysis.py +++ b/giant/analysis.py @@ -1706,16 +1706,27 @@ def plot_pdg_length_share(table: pl.DataFrame, max_slices: int = 6): # same event_ids can be overlaid against a real reference computed elsewhere. _ROLLOUT_COLS = [ - "event_id", "track_id", "termination_reason", - "pre_x", "pre_y", "pre_z", "pre_dx", "pre_dy", "pre_dz", "pre_E", - "post_x", "post_y", "post_z", "edep", + "event_id", + "track_id", + "termination_reason", + "pre_x", + "pre_y", + "pre_z", + "pre_dx", + "pre_dy", + "pre_dz", + "pre_E", + "post_x", + "post_y", + "post_z", + "edep", ] @dataclass class RolloutObservables: event_table: pd.DataFrame # one row per event_id (mm/MeV) - depth_edges: np.ndarray # (depth_bins+1,) mm, along shower axis + depth_edges: np.ndarray # (depth_bins+1,) mm, along shower axis transverse_edges: np.ndarray # (transverse_bins+1,) mm, perpendicular depth_profile: np.ndarray # (depth_bins,) mean edep/event/bin, MeV depth_profile_std: np.ndarray @@ -1736,9 +1747,7 @@ def compute_rollout_observables( per-event totals plus dataset-mean longitudinal/transverse profiles. """ pf = pq.ParquetFile(Path(path)) - coord = (pf.schema_arrow.metadata or {}).get( - PREDICT_COORD_METADATA_KEY.encode() - ) + coord = (pf.schema_arrow.metadata or {}).get(PREDICT_COORD_METADATA_KEY.encode()) if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE: raise ValueError( f"{path} is not a rollout file (coord={coord.decode()!r}); " @@ -1768,11 +1777,14 @@ def compute_rollout_observables( if not (depth_hi - depth_lo > 1e-6 * max(abs(depth_hi), 1.0)): depth_lo, depth_hi = depth_lo - 0.5, depth_hi + 0.5 depth_edges = np.linspace(depth_lo, depth_hi, depth_bins + 1) - transverse_edges = np.linspace(0.0, max(np.quantile(transverse, 0.999), 1e-6), - transverse_bins + 1) + transverse_edges = np.linspace( + 0.0, max(np.quantile(transverse, 0.999), 1e-6), transverse_bins + 1 + ) d_bin = np.clip(np.digitize(depth, depth_edges) - 1, 0, depth_bins - 1) - t_bin = np.clip(np.digitize(transverse, transverse_edges) - 1, 0, transverse_bins - 1) + t_bin = np.clip( + np.digitize(transverse, transverse_edges) - 1, 0, transverse_bins - 1 + ) # Per-(event, bin) edep sums, then mean/std across events. depth_ev = np.zeros((n_events, depth_bins)) @@ -1789,7 +1801,8 @@ def compute_rollout_observables( per_ev["n_tracks"] = df.groupby("event_id")["track_id"].nunique() per_ev["leaked_E"] = ( df.assign(_leak=np.where(is_leak, df["pre_E"], 0.0)) - .groupby("event_id")["_leak"].sum() + .groupby("event_id")["_leak"] + .sum() ) # Energy-weighted centroid depth per event (from the binned sums). bin_centers = 0.5 * (depth_edges[:-1] + depth_edges[1:]) @@ -1820,8 +1833,11 @@ def plot_rollout_longitudinal(obs: RolloutObservables, reference=None): fig, ax = plt.subplots(figsize=(7, 4)) ax.plot(centers, obs.depth_profile, label="rollout", color="C0") ax.fill_between( - centers, obs.depth_profile - obs.depth_profile_std, - obs.depth_profile + obs.depth_profile_std, alpha=0.2, color="C0", + centers, + obs.depth_profile - obs.depth_profile_std, + obs.depth_profile + obs.depth_profile_std, + alpha=0.2, + color="C0", ) if reference is not None: rc = 0.5 * (reference.depth_edges[:-1] + reference.depth_edges[1:]) @@ -1840,8 +1856,11 @@ def plot_rollout_transverse(obs: RolloutObservables, reference=None): fig, ax = plt.subplots(figsize=(7, 4)) ax.plot(centers, obs.transverse_profile, label="rollout", color="C0") ax.fill_between( - centers, obs.transverse_profile - obs.transverse_profile_std, - obs.transverse_profile + obs.transverse_profile_std, alpha=0.2, color="C0", + centers, + obs.transverse_profile - obs.transverse_profile_std, + obs.transverse_profile + obs.transverse_profile_std, + alpha=0.2, + color="C0", ) if reference is not None: rc = 0.5 * (reference.transverse_edges[:-1] + reference.transverse_edges[1:]) @@ -1858,11 +1877,22 @@ def plot_rollout_transverse(obs: RolloutObservables, reference=None): def plot_rollout_total_energy(obs: RolloutObservables, bins: int = 50, reference=None): """Distribution of total deposited energy per shower.""" fig, ax = plt.subplots(figsize=(7, 4)) - ax.hist(obs.event_table["total_edep"], bins=bins, histtype="step", - label="rollout", color="C0") + ax.hist( + obs.event_table["total_edep"], + bins=bins, + histtype="step", + label="rollout", + color="C0", + ) if reference is not None: - ax.hist(reference.event_table["real_total_edep"].to_numpy(), bins=bins, - histtype="step", label="real", color="k", ls="--") + ax.hist( + reference.event_table["real_total_edep"].to_numpy(), + bins=bins, + histtype="step", + label="real", + color="k", + ls="--", + ) ax.set_xlabel("total E_dep / event [MeV]") ax.set_ylabel("events") ax.set_title("Total deposited energy") diff --git a/giant/cli.py b/giant/cli.py index 4b61b10..d212bea 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -640,9 +640,7 @@ def predict( typer.echo(f"wrote {total:,} rows → {out}") -def _seed_from_data( - files: list[Path], n_events: int | None -) -> dict[str, np.ndarray]: +def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.ndarray]: """Pick each event's primary entry state (argmax-pre_E row) as a shower seed. Streams conditioning columns and keeps the highest-pre_E step per event_id — @@ -686,25 +684,30 @@ def rollout( Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)") ], checkpoint: Annotated[ - Path, typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)") + Path, + typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"), ], geometry: Annotated[ Path, typer.Option( - "--geometry", "-g", help="Geometry oracle .pkl (dwarf build-geometry-oracle)" + "--geometry", + "-g", + help="Geometry oracle .pkl (dwarf build-geometry-oracle)", ), ], energy_cutoff: Annotated[ float, typer.Option( - "--energy-cutoff", help="Stop a track when its energy drops below this [MeV]" + "--energy-cutoff", + help="Stop a track when its energy drops below this [MeV]", ), ] = 0.1, max_steps: Annotated[ int, typer.Option("--max-steps", help="Max steps per individual track") ] = 1000, steps: Annotated[ - int, typer.Option("--steps", "-s", help="Flow matching ODE steps per model call") + int, + typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"), ] = 10, batch_size: Annotated[ int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward") @@ -733,7 +736,8 @@ def rollout( Optional[Path], typer.Option("--out", "-o", help="Output steps parquet") ] = None, seed: Annotated[ - Optional[int], typer.Option("--seed", help="Torch/numpy seed for reproducibility") + Optional[int], + typer.Option("--seed", help="Torch/numpy seed for reproducibility"), ] = None, ) -> None: """Roll the surrogate forward into full showers (autoregressive).""" @@ -759,7 +763,9 @@ def rollout( cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) - model = DenoisingMLP(**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}) + model = DenoisingMLP( + **{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS} + ) model.load_state_dict(ckpt["model"]) model.to(_device).eval() sec_decoder = SecondaryDecoder( @@ -824,9 +830,7 @@ def rollout( ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False)) n_rows = len(records["event_id"]) - reasons = Counter( - r for r in records["termination_reason"].tolist() if r - ) + reasons = Counter(r for r in records["termination_reason"].tolist() if r) typer.echo(f"wrote {n_rows:,} step rows → {out}") typer.echo(f"terminations: {dict(reasons)}") typer.echo(f"reference: {ref_path}") diff --git a/giant/constants.py b/giant/constants.py index 230d1e3..01a16f7 100644 --- a/giant/constants.py +++ b/giant/constants.py @@ -13,7 +13,7 @@ K_MAX = 15 # EMB_DIM (continuous type embedding). EMB_DIM must match DenoisingMLP.emb_dim. # Default emb_dim=16 → SEC_SLOT_DIM=20. SEC_SLOT_DIM = 20 # 1 + 3 + 16 -EMB_DIM = 16 # must match model emb_dim default +EMB_DIM = 16 # must match model emb_dim default # Per-slot continuous (non-embedding) width: stick-breaking logit + local dir. CONT_SLOT_DIM = SEC_SLOT_DIM - EMB_DIM # 4 diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 433c5f3..e96b8f5 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -112,11 +112,22 @@ class StreamingStepsDataset(IterableDataset): buf_n += len(cond_cont) if buf_n >= self.shuffle_buffer: - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_n = ( - yield from self._flush( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, - final=False, - ) + ( + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + buf_n, + ) = yield from self._flush( + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + final=False, ) if buf_n > 0: @@ -164,7 +175,11 @@ class StreamingStepsDataset(IterableDataset): return [], [], [], [], [], [], 0 rem = n_full * bs return ( - [cont[rem:]], [cat[rem:]], [tgt[rem:]], - [nsec[rem:]], [sec[rem:]], [spdg[rem:]], + [cont[rem:]], + [cat[rem:]], + [tgt[rem:]], + [nsec[rem:]], + [sec[rem:]], + [spdg[rem:]], n - rem, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index fa0ba79..b51fd69 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -60,9 +60,7 @@ def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray: return out -def _pad_dir_col( - dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int -) -> np.ndarray: +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. diff --git a/giant/data/transforms.py b/giant/data/transforms.py index f5df8db..c068d63 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -485,7 +485,9 @@ def build_features( 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) - n_sec_raw = data["n_sec"].astype(np.int64) # (N,) unclamped, for the valid-slot mask + n_sec_raw = data["n_sec"].astype( + np.int64 + ) # (N,) unclamped, for the valid-slot mask # Clamp the classification label to K_MAX: the head only has K_MAX+1 classes # (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already # applied to sec_cont/sec_pdg_idx by the loader's list padding. Without this, diff --git a/giant/model/network.py b/giant/model/network.py index b64b48e..d341630 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -121,7 +121,7 @@ class DenoisingMLP(nn.Module): cond_cont: torch.Tensor, cond_cat: torch.Tensor, ) -> torch.Tensor: - t_emb = self.time_emb(t) # (B, time_dim) + t_emb = self.time_emb(t) # (B, time_dim) c_emb = self.cond_enc(cond_cont, cond_cat) # (B, cond_out_dim) cond = torch.cat([t_emb, c_emb], dim=-1) x = self.input_proj(x_t) @@ -176,9 +176,9 @@ class SecondaryConditionEncoder(nn.Module): cond_cat: torch.Tensor, stage1_out: torch.Tensor, ) -> torch.Tensor: - base = self.base(cond_cont, cond_cat) # (B, cond_out_dim) - s1 = self.stage1_proj(stage1_out).tanh() # (B, stage1_proj_dim) - return self.fuse(torch.cat([base, s1], dim=-1)) # (B, out_dim) + base = self.base(cond_cont, cond_cat) # (B, cond_out_dim) + s1 = self.stage1_proj(stage1_out).tanh() # (B, stage1_proj_dim) + return self.fuse(torch.cat([base, s1], dim=-1)) # (B, out_dim) class SecondaryDecoder(nn.Module): diff --git a/giant/pipeline.py b/giant/pipeline.py index 593455c..10b5746 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -110,7 +110,7 @@ def run_train_job( emb_dim = m.get("emb_dim", EMB_DIM) # SEC_SLOT_DIM must match constants (1 stick + 3 dir + emb_dim) assert SEC_SLOT_DIM == 1 + 3 + emb_dim, ( - f"SEC_SLOT_DIM={SEC_SLOT_DIM} must equal 1+3+emb_dim={1+3+emb_dim}; " + f"SEC_SLOT_DIM={SEC_SLOT_DIM} must equal 1+3+emb_dim={1 + 3 + emb_dim}; " "update giant/constants.py if emb_dim changed" ) diff --git a/giant/rollout.py b/giant/rollout.py index 9faa7af..a0fe071 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -39,10 +39,31 @@ from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx # Record columns produced per step / per terminal marker. _RECORD_KEYS = [ - "event_id", "track_id", "parent_id", "generation", "step_no", "pdg", - "pre_x", "pre_y", "pre_z", "pre_E", "pre_dx", "pre_dy", "pre_dz", - "post_x", "post_y", "post_z", "post_E", "post_dx", "post_dy", "post_dz", - "edep", "step_length", "material", "layer_id", "n_sec_pred", + "event_id", + "track_id", + "parent_id", + "generation", + "step_no", + "pdg", + "pre_x", + "pre_y", + "pre_z", + "pre_E", + "pre_dx", + "pre_dy", + "pre_dz", + "post_x", + "post_y", + "post_z", + "post_E", + "post_dx", + "post_dy", + "post_dz", + "edep", + "step_length", + "material", + "layer_id", + "n_sec_pred", "termination_reason", ] @@ -88,7 +109,12 @@ class _Recorder: if chunks: out[k] = np.concatenate(chunks, axis=0) else: - out[k] = np.empty(0, dtype=object if k in ("material", "termination_reason") else np.float64) + out[k] = np.empty( + 0, + dtype=object + if k in ("material", "termination_reason") + else np.float64, + ) return out @@ -136,13 +162,26 @@ def _terminal_rows(tr: dict[str, np.ndarray], sel: np.ndarray, reason: str, edep dir_ = tr["pre_dir"][sel] n = int(sel.sum()) return dict( - event_id=tr["event_id"][sel], track_id=tr["track_id"][sel], - parent_id=tr["parent_id"][sel], generation=tr["generation"][sel], - step_no=tr["step_in_track"][sel], pdg=tr["pdg"][sel], - pre_x=pos[:, 0], pre_y=pos[:, 1], pre_z=pos[:, 2], pre_E=tr["pre_E"][sel], - pre_dx=dir_[:, 0], pre_dy=dir_[:, 1], pre_dz=dir_[:, 2], - post_x=pos[:, 0], post_y=pos[:, 1], post_z=pos[:, 2], - post_E=np.zeros(n), post_dx=dir_[:, 0], post_dy=dir_[:, 1], post_dz=dir_[:, 2], + event_id=tr["event_id"][sel], + track_id=tr["track_id"][sel], + parent_id=tr["parent_id"][sel], + generation=tr["generation"][sel], + step_no=tr["step_in_track"][sel], + pdg=tr["pdg"][sel], + pre_x=pos[:, 0], + pre_y=pos[:, 1], + pre_z=pos[:, 2], + pre_E=tr["pre_E"][sel], + pre_dx=dir_[:, 0], + pre_dy=dir_[:, 1], + pre_dz=dir_[:, 2], + post_x=pos[:, 0], + post_y=pos[:, 1], + post_z=pos[:, 2], + post_E=np.zeros(n), + post_dx=dir_[:, 0], + post_dy=dir_[:, 1], + post_dz=dir_[:, 2], edep=np.asarray(edep, dtype=np.float64).reshape(n), step_length=np.zeros(n), material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel], @@ -182,8 +221,11 @@ def rollout( pdg_emb_weight = stage1_model.pdg_embedding_weight() frontier, counts = make_seed_frontier( - seeds["event_id"], seeds["pdg"], seeds["pre_pos"], - seeds["pre_E"], seeds["pre_dir"], + seeds["event_id"], + seeds["pdg"], + seeds["pre_pos"], + seeds["pre_E"], + seeds["pre_dir"], ) rec = _Recorder() @@ -191,14 +233,26 @@ def rollout( next_parts: list[dict[str, np.ndarray]] = [] n_total = len(frontier["event_id"]) for start in range(0, n_total, batch_size): - chunk = { - k: v[start : start + batch_size] for k, v in frontier.items() - } + chunk = {k: v[start : start + batch_size] for k, v in frontier.items()} next_parts.append( _step_chunk( - chunk, stage1_model, sec_decoder, oracle, cond_norm, tgt_norm, - pdg_map, mat_map, pdg_map_inv, pdg_emb_weight, rec, counts, - energy_cutoff, max_steps, steps, device, max_tracks_per_event, + chunk, + stage1_model, + sec_decoder, + oracle, + cond_norm, + tgt_norm, + pdg_map, + mat_map, + pdg_map_inv, + pdg_emb_weight, + rec, + counts, + energy_cutoff, + max_steps, + steps, + device, + max_tracks_per_event, ) ) frontier = _concat_frontiers(next_parts) @@ -207,9 +261,23 @@ def rollout( def _step_chunk( - tr, stage1_model, sec_decoder, oracle, cond_norm, tgt_norm, pdg_map, mat_map, - pdg_map_inv, pdg_emb_weight, rec, counts, energy_cutoff, max_steps, steps, - device, max_tracks_per_event, + tr, + stage1_model, + sec_decoder, + oracle, + cond_norm, + tgt_norm, + pdg_map, + mat_map, + pdg_map_inv, + pdg_emb_weight, + rec, + counts, + energy_cutoff, + max_steps, + steps, + device, + max_tracks_per_event, ) -> dict[str, np.ndarray]: """Advance one chunk of tracks by a single step; return the next frontier.""" n = len(tr["event_id"]) @@ -225,19 +293,33 @@ def _step_chunk( # --- Pre-step termination gates (in priority order; each track picks one) --- stop = np.zeros(n, dtype=bool) escaped_sel = escaped & ~stop - rec.add(**_terminal_rows(tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum())))) + rec.add( + **_terminal_rows( + tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum())) + ) + ) stop |= escaped_sel unknown_sel = ~known_pdg & ~stop - rec.add(**_terminal_rows(tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel])) + rec.add( + **_terminal_rows( + tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel] + ) + ) stop |= unknown_sel cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop - rec.add(**_terminal_rows(tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel])) + rec.add( + **_terminal_rows( + tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel] + ) + ) stop |= cutoff_sel maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop - rec.add(**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])) + rec.add( + **_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel]) + ) stop |= maxstep_sel active = ~stop @@ -250,8 +332,12 @@ def _step_chunk( # --- Build conditioning and run the two stages --- cond_dict = { - "pre_pos": tr["pre_pos"], "pre_E": tr["pre_E"], "pre_dir": tr["pre_dir"], - "layer_id": layer_id, "material": material, "pdg": tr["pdg"], + "pre_pos": tr["pre_pos"], + "pre_E": tr["pre_E"], + "pre_dir": tr["pre_dir"], + "layer_id": layer_id, + "material": material, + "pdg": tr["pdg"], } cond_cont, cond_cat = build_cond_features(cond_dict, pdg_map, mat_map, cond_norm) cc = torch.from_numpy(cond_cont).float().to(device) @@ -264,12 +350,18 @@ def _step_chunk( edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"]) post_dir_local = raw[:, 3:6].copy() - post_dir_local /= np.clip(np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None) + post_dir_local /= np.clip( + np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None + ) post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local) travel_dir_local = raw[:, 6:9].copy() - travel_dir_local /= np.clip(np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None) - post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local) + travel_dir_local /= np.clip( + np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None + ) + post_pos = reconstruct_post_pos( + tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local + ) n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64) @@ -279,8 +371,12 @@ def _step_chunk( ) sec_pdg_idx = snap_type_to_pdg_idx(sec_type_emb, pdg_emb_weight) sec_E, sec_dir_world, sec_pdg_code, sec_valid = decode_secondaries( - sec_cont.cpu().numpy(), sec_pdg_idx.cpu().numpy(), n_sec_np, - e_sec, tr["pre_dir"], pdg_map_inv, + sec_cont.cpu().numpy(), + sec_pdg_idx.cpu().numpy(), + n_sec_np, + e_sec, + tr["pre_dir"], + pdg_map_inv, ) edep = edep.astype(np.float64) @@ -288,8 +384,14 @@ def _step_chunk( # --- Spawn secondaries (with per-event track cap) --- new_tracks, dropped_edep = _spawn_secondaries( - tr, post_pos, sec_valid, sec_E, sec_dir_world, sec_pdg_code, - counts, max_tracks_per_event, + tr, + post_pos, + sec_valid, + sec_E, + sec_dir_world, + sec_pdg_code, + counts, + max_tracks_per_event, ) # Energy bookkeeping so each step conserves exactly (edep + carried + post_E # == pre_E): the primary lost `e_sec` to secondaries, but the decoded @@ -303,31 +405,58 @@ def _step_chunk( natural = post_E <= 0.0 reason = np.where(natural, TERM_NATURAL_END, "").astype(object) rec.add( - event_id=tr["event_id"], track_id=tr["track_id"], parent_id=tr["parent_id"], - generation=tr["generation"], step_no=tr["step_in_track"], pdg=tr["pdg"], - pre_x=tr["pre_pos"][:, 0], pre_y=tr["pre_pos"][:, 1], pre_z=tr["pre_pos"][:, 2], - pre_E=tr["pre_E"], pre_dx=tr["pre_dir"][:, 0], pre_dy=tr["pre_dir"][:, 1], + event_id=tr["event_id"], + track_id=tr["track_id"], + parent_id=tr["parent_id"], + generation=tr["generation"], + step_no=tr["step_in_track"], + pdg=tr["pdg"], + pre_x=tr["pre_pos"][:, 0], + pre_y=tr["pre_pos"][:, 1], + pre_z=tr["pre_pos"][:, 2], + pre_E=tr["pre_E"], + pre_dx=tr["pre_dir"][:, 0], + pre_dy=tr["pre_dir"][:, 1], pre_dz=tr["pre_dir"][:, 2], - post_x=post_pos[:, 0], post_y=post_pos[:, 1], post_z=post_pos[:, 2], - post_E=post_E, post_dx=post_dir_world[:, 0], post_dy=post_dir_world[:, 1], - post_dz=post_dir_world[:, 2], edep=edep, step_length=step_length, - material=material, layer_id=layer_id, n_sec_pred=n_sec_np, + post_x=post_pos[:, 0], + post_y=post_pos[:, 1], + post_z=post_pos[:, 2], + post_E=post_E, + post_dx=post_dir_world[:, 0], + post_dy=post_dir_world[:, 1], + post_dz=post_dir_world[:, 2], + edep=edep, + step_length=step_length, + material=material, + layer_id=layer_id, + n_sec_pred=n_sec_np, termination_reason=reason, ) # --- Continue surviving primaries --- cont = ~natural cont_frontier = { - "event_id": tr["event_id"][cont], "track_id": tr["track_id"][cont], - "parent_id": tr["parent_id"][cont], "generation": tr["generation"][cont], - "step_in_track": tr["step_in_track"][cont] + 1, "pdg": tr["pdg"][cont], - "pre_pos": post_pos[cont], "pre_E": post_E[cont], "pre_dir": post_dir_world[cont], + "event_id": tr["event_id"][cont], + "track_id": tr["track_id"][cont], + "parent_id": tr["parent_id"][cont], + "generation": tr["generation"][cont], + "step_in_track": tr["step_in_track"][cont] + 1, + "pdg": tr["pdg"][cont], + "pre_pos": post_pos[cont], + "pre_E": post_E[cont], + "pre_dir": post_dir_world[cont], } return _concat_frontiers([cont_frontier, new_tracks]) def _spawn_secondaries( - tr, post_pos, sec_valid, sec_E, sec_dir_world, sec_pdg_code, counts, + tr, + post_pos, + sec_valid, + sec_E, + sec_dir_world, + sec_pdg_code, + counts, max_tracks_per_event, ) -> tuple[dict[str, np.ndarray], np.ndarray]: """Turn valid secondaries into new tracks; return (frontier, per-parent dropped edep). diff --git a/giant/sample.py b/giant/sample.py index c2480a1..900b32f 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -63,8 +63,8 @@ def sample_secondaries( sec_cont = x_slots[:, :, :4] sec_type_emb = x_slots[:, :, 4:] - sec_valid = ( - torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1) + sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze( + 1 ) return sec_cont, sec_type_emb, sec_valid diff --git a/tests/test_phase2.py b/tests/test_phase2.py index e8b4f2d..8e46740 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -1,4 +1,5 @@ """Tests for Phase 2: secondary particle prediction.""" + import numpy as np import pytest import torch @@ -11,6 +12,7 @@ from giant.sample import sample_secondaries, snap_type_to_pdg_idx # ── helpers ────────────────────────────────────────────────────────────────── + def _stage1(pdg=3, mat=2): return DenoisingMLP(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2) @@ -29,6 +31,7 @@ def _cond(B=8, pdg=3, mat=2): # ── DenoisingMLP Phase-2 additions ─────────────────────────────────────────── + def test_predict_n_sec_shape(): B = 8 model = _stage1() @@ -53,6 +56,7 @@ def test_pdg_embedding_weight_shape(): # ── SecondaryDecoder ────────────────────────────────────────────────────────── + def test_sec_decoder_output_shape(): B = 8 decoder = _sec_decoder() @@ -89,6 +93,7 @@ def test_sec_decoder_gradients(): # ── masked flow matching loss ───────────────────────────────────────────────── + def test_flow_matching_loss_secondary_scalar(): B, pdg, mat = 8, 3, 2 decoder = _sec_decoder(pdg, mat) @@ -96,7 +101,9 @@ def test_flow_matching_loss_secondary_scalar(): cond_cont, cond_cat = _cond(B, pdg, mat) stage1_out = torch.randn(B, X_DIM) sec_mask = torch.ones(B, K_MAX, dtype=torch.bool) - loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask) + loss = flow_matching_loss_secondary( + decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask + ) assert loss.shape == () assert loss.item() >= 0.0 @@ -109,7 +116,9 @@ def test_flow_matching_loss_secondary_mask_zeros_padding(): cond_cont, cond_cat = _cond(B, pdg, mat) stage1_out = torch.randn(B, X_DIM) sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool) - loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask) + loss = flow_matching_loss_secondary( + decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask + ) assert loss.item() == pytest.approx(0.0, abs=1e-6) @@ -128,6 +137,7 @@ def test_flow_matching_loss_secondary_has_grad(): # ── sampling ────────────────────────────────────────────────────────────────── + def test_sample_secondaries_shapes(): B, pdg, mat = 6, 3, 2 decoder = _sec_decoder(pdg, mat) @@ -169,6 +179,7 @@ def test_snap_type_to_pdg_idx_shape(): # ── encode_secondaries round-trip ───────────────────────────────────────────── + def test_encode_secondaries_energy_conservation(): """Decoded stick-breaking fractions must sum to ≈ e_sec.""" from giant.data.transforms import encode_secondaries @@ -217,6 +228,6 @@ def test_encode_secondaries_direction_encoding(): sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) # dir columns are sec_cont[:, :, 1:4] - local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only + local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only norms_out = np.linalg.norm(local_dirs, axis=-1) np.testing.assert_allclose(norms_out, 1.0, atol=1e-5) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 4674beb..f744755 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -56,16 +56,31 @@ def _seeds(n=6): } -def _run(escape_threshold=1e9, energy_cutoff=1.0, max_steps=30, - max_tracks_per_event=300, seeds=None): +def _run( + escape_threshold=1e9, + energy_cutoff=1.0, + max_steps=30, + max_tracks_per_event=300, + seeds=None, +): torch.manual_seed(0) np.random.seed(0) s1, s2 = _models() cond, tgt = _norms() return rollout( - s1, s2, _oracle(), seeds or _seeds(), cond, tgt, PDG_MAP, MAT_MAP, - energy_cutoff=energy_cutoff, max_steps=max_steps, steps=4, - batch_size=128, max_tracks_per_event=max_tracks_per_event, + s1, + s2, + _oracle(), + seeds or _seeds(), + cond, + tgt, + PDG_MAP, + MAT_MAP, + energy_cutoff=energy_cutoff, + max_steps=max_steps, + steps=4, + batch_size=128, + max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, )