lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Apply ruff format across the codebase

Browse Source 
Whitespace-only reflow (line wrapping, blank lines between defs); no
logic changes.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-08 14:44:53 +02:00
parent af8dce53a7
commit a14a4f973a
16 changed files with 351 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
analysis/compare_ode_steps_energy_conservation.py
+30 -11
View File
@@ -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")
analysis/compare_ode_steps_kl.py
+1 -3
View File
@@ -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}")
analysis/export_energy_conservation_poc.py
+6 -2
View File
@@ -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))
analysis/export_rollout_observables.py
+3 -1
View File
@@ -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} ===")
giant/analysis.py
+49 -19
View File
@@ -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")
giant/cli.py
+16 -12
View File
@@ -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}")
giant/constants.py
+1 -1
View File
@@ -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
giant/data/dataset.py
+22 -7
View File
@@ -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,
)
giant/data/loader.py
+1 -3
View File
@@ -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.
giant/data/transforms.py
+3 -1
View File
@@ -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,
giant/model/network.py
+4 -4
View File
@@ -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):
giant/pipeline.py
+1 -1
View File
@@ -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"
)
giant/rollout.py
+178 -49
View File
@@ -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).
giant/sample.py
+2 -2
View File
@@ -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
tests/test_phase2.py
+14 -3
View File
@@ -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)
tests/test_rollout.py
+20 -5
View File
@@ -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,
)