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

Bump ruff line-length to 120 and reformat
CI / Lint (ruff check) (push) Successful in 31s
Details
CI / Format (ruff format) (push) Successful in 32s
Details
CI / Sync project version with tag (push) Has been skipped
Details
CI / Lint (ruff check) (pull_request) Successful in 26s
Details
CI / Type check (ty) (push) Successful in 29s
Details
CI / Format (ruff format) (pull_request) Successful in 33s
Details
CI / Sync project version with tag (pull_request) Has been skipped
Details
CI / Type check (ty) (pull_request) Successful in 35s
Details
CI / Tests (pull_request) Successful in 3m47s
Details
CI / Tests (push) Successful in 3m55s
Details

Browse Source 
Rejoins lines that only wrapped because they exceeded the old 88-char
limit; ruff check and the full test suite (725 passed) are unaffected.
This commit is contained in:
Lars Bogner
2026-08-12 13:33:09 +02:00
parent 9ce7b32324
commit 55332db67a
66 changed files with 757 additions and 2413 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/analysis/catalog.py
+21 -61
View File
@@ -166,9 +166,7 @@ def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]:
return list(merged.get(str(key), [0] * nbins))
def _np_hist_pair(
r: np.ndarray, t: np.ndarray, nbins: int
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
def _np_hist_pair(r: np.ndarray, t: np.ndarray, nbins: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Shared-edge histogram of two small per-event arrays (robust range)."""
both = np.concatenate([r, t]) if (len(r) or len(t)) else np.array([0.0, 1.0])
lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
@@ -240,9 +238,7 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
ids, bins = event_energy_bins(lf, edges)
return pl.col("event_id").replace_strict(
ids, bins, default=-1, return_dtype=pl.Int64
)
return pl.col("event_id").replace_strict(ids, bins, default=-1, return_dtype=pl.Int64)
def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
@@ -265,9 +261,7 @@ def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
}
def _marginal_grouped_finalize(
parts: list[dict], ctx: Context, var: str, axis: str
) -> Reduced:
def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis: str) -> Reduced:
label, _ = _var(var)
edges = _marginal_edges(ctx, var)
nb = len(edges) - 1
@@ -317,9 +311,7 @@ def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
return {"r": r.tolist(), "t": t.tolist()}
def _event_scalar_finalize(
parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str
) -> Reduced:
def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str) -> Reduced:
r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
edges, rc, tc = _np_hist_pair(r, t, ctx.n_marginal_bins)
@@ -488,9 +480,7 @@ def _leakage_partial(b: Bundle) -> dict:
def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts])
edges = np.linspace(
0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1
)
edges = np.linspace(0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1)
counts = np.histogram(frac, edges)[0]
return Reduced(
id="leakage_fraction",
@@ -521,18 +511,8 @@ def _sec_frames(b: Bundle):
def _sec_count_per_event_partial(b: Bundle) -> dict:
r_sec, t_sec = _sec_frames(b)
r = (
r_sec.group_by("event_id")
.agg(pl.len().alias("n"))
.collect(engine="streaming")["n"]
.to_numpy()
)
t = (
t_sec.group_by("event_id")
.agg(pl.len().alias("n"))
.collect(engine="streaming")["n"]
.to_numpy()
)
r = r_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
t = t_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
return {"r": r.tolist(), "t": t.tolist()}
@@ -568,9 +548,7 @@ def _sec_count_per_species_partial(b: Bundle) -> dict:
def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
r = sum_merge([p["r"] for p in parts])
t = sum_merge([p["t"] for p in parts])
keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[
: len(ctx.top_pdgs)
]
keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[: len(ctx.top_pdgs)]
return Reduced(
id="sec_count_per_species",
family="secondaries",
@@ -617,11 +595,9 @@ def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
def _sec_cos_angle_partial(b: Bundle) -> dict:
edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1)
cos = (
pl.col("sdx") * pl.col("axis_x")
+ pl.col("sdy") * pl.col("axis_y")
+ pl.col("sdz") * pl.col("axis_z")
).clip(-1.0, 1.0)
cos = (pl.col("sdx") * pl.col("axis_x") + pl.col("sdy") * pl.col("axis_y") + pl.col("sdz") * pl.col("axis_z")).clip(
-1.0, 1.0
)
def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]:
ea = entry_axis(steps_lf)
@@ -659,15 +635,13 @@ _router_gating_partial, _router_gating_finalize = _unchunkable(
lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys)
)
_router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
lambda b: compute_router_share_by_pdg(
b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs
)
lambda b: compute_router_share_by_pdg(b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs)
)
_router_share_process_partial, _router_share_process_finalize = _unchunkable(
lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
)
_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = (
_unchunkable(lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist))
_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = _unchunkable(
lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist)
)
@@ -689,9 +663,7 @@ def build_catalog() -> list[PlotSpec]:
f"marginal_{var}",
"marginals",
compute_partial=lambda b, v=var: _marginal_overall_partial(b, v),
finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(
parts, ctx, v
),
finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(parts, ctx, v),
)
)
for axis in GROUPING_AXES:
@@ -699,12 +671,8 @@ def build_catalog() -> list[PlotSpec]:
PlotSpec(
f"marginal_{var}_by_{axis}",
"marginals",
compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(
b, v, a
),
finalize=lambda parts, ctx, v=var, a=axis: (
_marginal_grouped_finalize(parts, ctx, v, a)
),
compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(b, v, a),
finalize=lambda parts, ctx, v=var, a=axis: _marginal_grouped_finalize(parts, ctx, v, a),
)
)
@@ -712,9 +680,7 @@ def build_catalog() -> list[PlotSpec]:
PlotSpec(
"event_total_edep",
"event",
compute_partial=lambda b: _event_scalar_partial(
b, "total_edep", use_all=True
),
compute_partial=lambda b: _event_scalar_partial(b, "total_edep", use_all=True),
finalize=lambda parts, ctx: _event_scalar_finalize(
parts,
ctx,
@@ -732,9 +698,7 @@ def build_catalog() -> list[PlotSpec]:
PlotSpec(
"event_mean_length",
"event",
compute_partial=lambda b: _event_scalar_partial(
b, "mean_length", use_all=False
),
compute_partial=lambda b: _event_scalar_partial(b, "mean_length", use_all=False),
finalize=lambda parts, ctx: _event_scalar_finalize(
parts,
ctx,
@@ -746,9 +710,7 @@ def build_catalog() -> list[PlotSpec]:
PlotSpec(
"event_n_steps",
"event",
compute_partial=lambda b: _event_scalar_partial(
b, "n_steps", use_all=False
),
compute_partial=lambda b: _event_scalar_partial(b, "n_steps", use_all=False),
finalize=lambda parts, ctx: _event_scalar_finalize(
parts,
ctx,
@@ -773,9 +735,7 @@ def build_catalog() -> list[PlotSpec]:
PlotSpec(
"shower_transverse",
"shower",
compute_partial=lambda b: _profile_partial(
b, transverse_expr, "transverse_edges"
),
compute_partial=lambda b: _profile_partial(b, transverse_expr, "transverse_edges"),
finalize=lambda parts, ctx: _profile_finalize(
parts,
ctx,
giant/analysis/condor.py
+6 -20
View File
@@ -161,9 +161,7 @@ class RunMeta:
return cls(**json.loads(Path(path).read_text()))
def _rows_per_chunk(
rollout: str | Path, reference: str | Path, n_chunks: int
) -> list[int]:
def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -> list[int]:
"""Rollout+reference row count of each ``event_id % n_chunks`` chunk.
One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
@@ -268,8 +266,7 @@ def compute_reduced(
effective_n = n_chunks if spec.chunkable else 1
if not (0 <= chunk_index < effective_n):
raise ValueError(
f"{spec_id}: chunk_index={chunk_index} out of range for "
f"n_chunks={effective_n} (chunkable={spec.chunkable})"
f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})"
)
bundle = Bundle.open(
rollout,
@@ -327,10 +324,7 @@ def merge_one(spec_id: str, run_dir: str | Path) -> Path:
effective_n = meta.n_chunks if spec.chunkable else 1
partial_dir = run_path / "reduced_partial"
found = {
p.chunk: p
for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))
}
found = {p.chunk: p for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))}
missing = sorted(set(range(effective_n)) - set(found))
if missing:
raise FileNotFoundError(
@@ -375,11 +369,7 @@ exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
reqs_attrs = (
"+RemoteJob = True\n"
if cfg.remote
else "requirements = TARGET.ProvidesETPResources\n"
)
reqs_attrs = "+RemoteJob = True\n" if cfg.remote else "requirements = TARGET.ProvidesETPResources\n"
return (
"universe = docker\n"
f"docker_image = {cfg.docker_image}\n"
@@ -399,9 +389,7 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> st
)
def _job_walltimes(
run_dir: Path, ids: list[str], n_chunks: int
) -> list[tuple[str, int, int]]:
def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]:
"""``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.
Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
@@ -477,9 +465,7 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
(run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)
wrapper = run_dir / "run_compute.sh"
wrapper.write_text(
_WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir)
)
wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir))
wrapper.chmod(0o755)
jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
giant/analysis/context.py
+6 -25
View File
@@ -74,9 +74,7 @@ def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFram
return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
def _combined_quantiles(
r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
) -> tuple[float, float]:
def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float) -> tuple[float, float]:
"""Robust (lo_q, hi_q) range over the union of two value samples."""
both = np.concatenate([r_vals, t_vals])
lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
@@ -104,31 +102,15 @@ def build_context(
# Ranged marginal variables: robust ranges over a shared row subsample.
exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
r_s = (
_row_subsample(r_lf, sample_rows, seed)
.select(exprs)
.collect(engine="streaming")
)
t_s = (
_row_subsample(t_lf, sample_rows, seed)
.select(exprs)
.collect(engine="streaming")
)
r_s = _row_subsample(r_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
t_s = _row_subsample(t_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
var_ranges = {
name: _combined_quantiles(
r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
)
for name in RANGED_VARS
name: _combined_quantiles(r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q) for name in RANGED_VARS
}
# Energy-bin edges from exact per-event incident energies (cheap group_by).
def _incident(lf: pl.LazyFrame) -> np.ndarray:
return (
lf.group_by("event_id")
.agg(pl.col("pre_E").max())
.collect(engine="streaming")["pre_E"]
.to_numpy()
)
return lf.group_by("event_id").agg(pl.col("pre_E").max()).collect(engine="streaming")["pre_E"].to_numpy()
r_inc, t_inc = _incident(r_lf), _incident(t_lf)
energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
@@ -145,8 +127,7 @@ def build_context(
)
top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
materials = sorted(
set(_counts(r_lf, "material")["material"].to_list())
| set(_counts(t_lf, "material")["material"].to_list())
set(_counts(r_lf, "material")["material"].to_list()) | set(_counts(t_lf, "material")["material"].to_list())
)
# Shower depth / transverse ranges from a subsampled proxy.
giant/analysis/grouping.py
+2 -6
View File
@@ -68,9 +68,7 @@ def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
def energy_bin_labels(edges: np.ndarray) -> list[str]:
"""``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
return [
f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
]
return [f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)]
def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
@@ -88,9 +86,7 @@ def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
return idx.clip(0, n_bins - 1)
def event_energy_bins(
lf: pl.LazyFrame, edges: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
def event_energy_bins(lf: pl.LazyFrame, edges: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.
Incident energy is ``max(pre_E)`` per event (the primary). One bounded
giant/analysis/reduce.py
+2 -7
View File
@@ -142,9 +142,7 @@ def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
"""
ids = entry["event_id"].to_numpy()
return lf.with_columns(
pl.col("event_id")
.replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64)
.alias(col)
pl.col("event_id").replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64).alias(col)
for col in _ENTRY_AXIS_COLS
)
@@ -254,10 +252,7 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
lf.group_by("event_id")
.agg(
pl.col("edep").sum().alias("deposited"),
pl.col("pre_E")
.filter(pl.col("termination_reason") == TERM_ESCAPED)
.sum()
.alias("escaped"),
pl.col("pre_E").filter(pl.col("termination_reason") == TERM_ESCAPED).sum().alias("escaped"),
)
.collect(engine="streaming")
)
giant/analysis/render.py
+8 -24
View File
@@ -139,9 +139,7 @@ def _render_overlay(r: Reduced, params: dict):
def _render_single(r: Reduced, params: dict):
edges = np.asarray(r.payload["edges"])
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
ax.stairs(
_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
)
ax.stairs(_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"])
if r.payload.get("log_y"):
ax.set_yscale("log")
if r.payload.get("log_x"):
@@ -187,9 +185,7 @@ def _render_profile(r: Reduced, params: dict):
mean = np.asarray(r.payload[f"{key}_mean"])
std = np.asarray(r.payload[f"{key}_std"])
(line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
ax.fill_between(
centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
)
ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=line.get_color())
ax.set_xlabel(r.xlabel)
ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
ps.style_legend(ax, title="source")
@@ -201,9 +197,7 @@ def _render_bar(r: Reduced, params: dict):
x = np.arange(len(labels))
width = 0.4
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
ax.bar(
x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
)
ax.bar(x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"])
ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
@@ -215,9 +209,7 @@ def _render_bar(r: Reduced, params: dict):
def _render_router_gating(r: Reduced, params: dict):
n_experts = r.payload["n_experts"]
log_x = r.payload.get("log_x", False)
fig, axes = ps.new_figure(
"slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
)
fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False)
flat = axes.ravel()
for ax, key in zip(flat, ("rollout", "reference")):
side = r.payload.get(key, {})
@@ -226,9 +218,7 @@ def _render_router_gating(r: Reduced, params: dict):
if len(centers) and means.size:
cum = np.zeros(len(centers))
for i in range(n_experts):
ax.fill_between(
centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
)
ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}")
cum = cum + means[:, i]
if log_x:
ax.set_xscale("log")
@@ -347,9 +337,7 @@ def render_all(
families.add(r.family)
fig = render(r, run_meta)
ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
(family_dir / f"{r.id}.yaml").write_text(
yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
)
(family_dir / f"{r.id}.yaml").write_text(yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False))
pdfs.append(family_dir / f"{r.id}.pdf")
import matplotlib.pyplot as plt
@@ -370,9 +358,7 @@ def render_all(
)
for fam in families:
(out_dir / fam / "metadata.yaml").write_text(
yaml.safe_dump(
{"title": fam, "description": f"{fam} plots."}, sort_keys=False
)
yaml.safe_dump({"title": fam, "description": f"{fam} plots."}, sort_keys=False)
)
if run_gallery:
@@ -400,6 +386,4 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
"reference": meta.reference,
**meta.plot_meta,
}
return render_all(
run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
)
return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)
giant/analysis/router_gating.py
+11 -34
View File
@@ -95,9 +95,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
# New nested shape (has a "stage1_model" key) vs. a v0.2 checkpoint's
# flat model_config.
router_cfg = (
(model_cfg.get("stage1_model") or {}).get("router")
if "stage1_model" in model_cfg
else model_cfg.get("router")
(model_cfg.get("stage1_model") or {}).get("router") if "stage1_model" in model_cfg else model_cfg.get("router")
)
if not router_cfg or not router_cfg.get("enabled"):
return None
@@ -124,9 +122,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
)
def _subsample(
lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()
) -> pl.DataFrame:
def _subsample(lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()) -> pl.DataFrame:
total = lf.select(pl.len()).collect(engine="streaming").item()
if total > n:
threshold = int(n / total * 2**32)
@@ -134,9 +130,7 @@ def _subsample(
return lf.select(*_COLS, *extra_cols).collect(engine="streaming")
def _gate_for_df(
handle: _RouterHandle, df: pl.DataFrame
) -> tuple[pl.DataFrame, np.ndarray]:
def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame, np.ndarray]:
"""(filtered df, gate_weights) for rows in ``df`` with a known pdg/material.
Rows whose species or material never appeared in the checkpoint's
@@ -160,13 +154,9 @@ def _gate_for_df(
df = df.filter(pl.Series(known, dtype=pl.Boolean))
data = {
"pre_pos": np.column_stack(
[df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]
),
"pre_pos": np.column_stack([df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]),
"pre_E": df["pre_E"].to_numpy(),
"pre_dir": np.column_stack(
[df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]
),
"pre_dir": np.column_stack([df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]),
"layer_id": df["layer_id"].to_numpy(),
"pdg": df["pdg"].to_numpy(),
"material": df["material"].to_numpy(),
@@ -180,9 +170,7 @@ def _gate_for_df(
material_conditioning=handle.material_conditioning,
)
with torch.no_grad():
gate = handle.router.gate(
torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()
).numpy()
gate = handle.router.gate(torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()).numpy()
return df, gate
@@ -205,9 +193,7 @@ def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict:
return {"centers": centers[valid].tolist(), "means": means[valid].tolist()}
def _top1_shares(
categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int
) -> dict[str, list[float]]:
def _top1_shares(categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int) -> dict[str, list[float]]:
"""Fraction of each category's rows hard-dispatched to each expert.
Uses `Router.top1` (argmax), not the soft `gate` mean — grouped top-1
@@ -227,10 +213,7 @@ def _top1_shares(
return shares
_NOTE_NOT_MOE = (
"checkpoint has no enabled MoE router (model.router.enabled is "
"false/absent) — nothing to show"
)
_NOTE_NOT_MOE = "checkpoint has no enabled MoE router (model.router.enabled is false/absent) — nothing to show"
_TITLES = {
"router_gating": "Router gating (mixture-of-experts decision boundaries)",
@@ -266,9 +249,7 @@ def compute_router_gating(
df = _subsample(lf, _SAMPLE_ROWS, seed)
df, gate = _gate_for_df(handle, df)
x = df["pre_E"].to_numpy()
sides[name] = (
_quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
)
sides[name] = _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
return Reduced(
id="router_gating",
@@ -304,9 +285,7 @@ def compute_router_share_by_pdg(
df, gate = _gate_for_df(handle, df)
if len(df):
idx = gate.argmax(axis=1)
shares = _top1_shares(
df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts
)
shares = _top1_shares(df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts)
else:
shares = {str(p): [0.0] * handle.router.n_experts for p in top_pdgs}
sides[name] = {labels[i]: shares[str(p)] for i, p in enumerate(top_pdgs)}
@@ -351,9 +330,7 @@ def compute_router_share_by_process(
counts = df["process"].value_counts().sort("count", descending=True)
order = counts["process"].to_list()[:top_k]
idx = gate.argmax(axis=1)
shares = _top1_shares(
df["process"].to_numpy(), idx, order, handle.router.n_experts
)
shares = _top1_shares(df["process"].to_numpy(), idx, order, handle.router.n_experts)
else:
order, shares = [], {}
giant/analysis/runtime_estimate.py
+1 -3
View File
@@ -54,9 +54,7 @@ _FIXED_OVERHEAD_S = 60.0
# scan. Calibrated from the 3 real router jobs' observed wall times (119, 66,
# 124s) — max minus _FIXED_OVERHEAD_S, on top of it.
_ROUTER_FIXED_S = 64.0
_ROUTER_IDS = frozenset(
{"router_gating", "router_share_by_pdg", "router_share_by_process"}
)
_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process"})
# Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot
# added after the last calibration run) — the most expensive fitted per-row
giant/analysis/sources.py
+3 -12
View File
@@ -93,10 +93,7 @@ def _check_rollout_metadata(path: Path) -> None:
metadata = pq.read_schema(path).metadata or {}
coord = metadata.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}); "
"expected `giant rollout` output"
)
raise ValueError(f"{path} is not a rollout file (coord={coord.decode()!r}); expected `giant rollout` output")
def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
@@ -121,11 +118,7 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
else:
# The reference (a rollout's seed `dataset`) may be a directory of
# parquet shards rather than a single file — scan them all.
lf = (
pl.scan_parquet(str(path / "**/*.parquet"))
if path.is_dir()
else pl.scan_parquet(path)
)
lf = pl.scan_parquet(str(path / "**/*.parquet")) if path.is_dir() else pl.scan_parquet(path)
return lf.with_columns(pl.col("pdg").cast(pl.Int64))
@@ -137,9 +130,7 @@ def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
"""
if side is Side.reference:
return lf
return lf.filter(
~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS))
)
return lf.filter(~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS)))
def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
giant/cli.py
+97 -261
View File
@@ -65,9 +65,7 @@ def _router_total_experts(router_cfg: dict) -> int:
"""
if router_cfg.get("type") == "composed":
axis_counts = {
m.group(1): int(v)
for k, v in router_cfg.items()
if (m := re.match(r"^axis(\d+)_n_experts$", k))
m.group(1): int(v) for k, v in router_cfg.items() if (m := re.match(r"^axis(\d+)_n_experts$", k))
}
return math.prod(axis_counts.values()) if axis_counts else 1
return int(router_cfg.get("n_experts", 1))
@@ -104,11 +102,7 @@ def _ddpm_steps(model_cfg: dict, stage: str) -> int:
def _particle_type_other_policy(model_cfg: dict) -> str:
return (
_stage_cfg(model_cfg, "stage2")
.get("particle_type", {})
.get("other_policy", "sample")
)
return _stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")
def _load_pdg_topn_map(ckpt: dict):
@@ -129,9 +123,7 @@ def _load_mat_topn_map(ckpt: dict):
return topnmap_from_json(raw, axis="material") if raw is not None else None
def _batch_size_estimate_dims(
model_cfg: dict, training: bool, stage: str = "stage1"
) -> tuple[int, int]:
def _batch_size_estimate_dims(model_cfg: dict, training: bool, stage: str = "stage1") -> tuple[int, int]:
"""Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.
`model_cfg` is either the new nested shape (has a `f"{stage}_model"` key
@@ -336,14 +328,10 @@ def _load_model_weights(
@app.command()
def train(
data: Annotated[
Path, typer.Argument(help="Parquet file or directory of parquet files")
],
data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
config: Annotated[
Optional[Path],
typer.Option(
"--config", "-c", help="TOML config file (overridden by explicit flags)"
),
typer.Option("--config", "-c", help="TOML config file (overridden by explicit flags)"),
] = None,
mode: Annotated[
Optional[Mode],
@@ -355,8 +343,7 @@ def train(
typer.Option(
"--batch-size",
"-b",
help="Integer, or 'auto' to estimate from free GPU memory "
"(cuda devices only)",
help="Integer, or 'auto' to estimate from free GPU memory (cuda devices only)",
),
] = None,
lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
@@ -372,17 +359,13 @@ def train(
"alongside the raw weights in checkpoints (0 disables; default: 0.9999)",
),
] = None,
warmup_epochs: Annotated[
Optional[int], typer.Option("--warmup-epochs", "-w")
] = None,
warmup_epochs: Annotated[Optional[int], typer.Option("--warmup-epochs", "-w")] = None,
hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
dropout: Annotated[
Optional[float],
typer.Option(
"--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
),
typer.Option("--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"),
] = None,
stage1_generator: Annotated[
Optional[Mode],
@@ -402,9 +385,7 @@ def train(
] = None,
stage1_n_res_blocks: Annotated[
Optional[int],
typer.Option(
"--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"
),
typer.Option("--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"),
] = None,
stage1_dropout: Annotated[
Optional[float],
@@ -452,8 +433,7 @@ def train(
Optional[int],
typer.Option(
"--stage2-context-dim",
help="Width of the projected stage-1 outcome fed into stage 2's "
"conditioning (default: 64)",
help="Width of the projected stage-1 outcome fed into stage 2's conditioning (default: 64)",
),
] = None,
stage2_stage1_context: Annotated[
@@ -485,13 +465,9 @@ def train(
] = None,
router_type: Annotated[
Optional[str],
typer.Option(
"--router-type", help="Router implementation name (see ROUTER_REGISTRY)"
),
] = None,
n_experts: Annotated[
Optional[int], typer.Option("--n-experts", help="Number of routed experts")
typer.Option("--router-type", help="Router implementation name (see ROUTER_REGISTRY)"),
] = None,
n_experts: Annotated[Optional[int], typer.Option("--n-experts", help="Number of routed experts")] = None,
router_axis: Annotated[
Optional[list[str]],
typer.Option(
@@ -506,32 +482,28 @@ def train(
Optional[int],
typer.Option(
"--n-critic",
help="WGAN-GP (--mode wgan only): critic updates per generator "
"update (default: 5)",
help="WGAN-GP (--mode wgan only): critic updates per generator update (default: 5)",
),
] = None,
gp_weight: Annotated[
Optional[float],
typer.Option(
"--gp-weight",
help="WGAN-GP (--mode wgan only): gradient-penalty coefficient "
"(default: 10.0)",
help="WGAN-GP (--mode wgan only): gradient-penalty coefficient (default: 10.0)",
),
] = None,
noise_dim: Annotated[
Optional[int],
typer.Option(
"--noise-dim",
help="WGAN (--mode wgan only): generator input noise-vector "
"width (default: 64)",
help="WGAN (--mode wgan only): generator input noise-vector width (default: 64)",
),
] = None,
critic_lr: Annotated[
Optional[float],
typer.Option(
"--critic-lr",
help="WGAN-GP (--mode wgan only): critic learning rate "
"(default: same as --lr)",
help="WGAN-GP (--mode wgan only): critic learning rate (default: same as --lr)",
),
] = None,
stage1_n_critic: Annotated[
@@ -540,21 +512,15 @@ def train(
] = None,
stage1_gp_weight: Annotated[
Optional[float],
typer.Option(
"--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"
),
typer.Option("--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"),
] = None,
stage1_noise_dim: Annotated[
Optional[int],
typer.Option(
"--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"
),
typer.Option("--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"),
] = None,
stage1_critic_lr: Annotated[
Optional[float],
typer.Option(
"--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"
),
typer.Option("--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"),
] = None,
stage2_n_critic: Annotated[
Optional[int],
@@ -562,9 +528,7 @@ def train(
] = None,
stage2_gp_weight: Annotated[
Optional[float],
typer.Option(
"--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"
),
typer.Option("--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"),
] = None,
stage2_noise_dim: Annotated[
Optional[int],
@@ -576,13 +540,9 @@ def train(
] = None,
stage2_critic_lr: Annotated[
Optional[float],
typer.Option(
"--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"
),
] = None,
val_fraction: Annotated[
Optional[float], typer.Option("--val-fraction", "-f")
typer.Option("--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"),
] = None,
val_fraction: Annotated[Optional[float], typer.Option("--val-fraction", "-f")] = None,
seed: Annotated[
Optional[int],
typer.Option("--seed", "-s", help="Random seed for reproducibility"),
@@ -608,15 +568,12 @@ def train(
Optional[int],
typer.Option(
"--max-val-batches",
help="Cap the per-epoch val-loss pass to N batches (0 = full "
"val set every epoch; default: 200)",
help="Cap the per-epoch val-loss pass to N batches (0 = full val set every epoch; default: 200)",
),
] = None,
shuffle_buffer: Annotated[
int,
typer.Option(
"--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"
),
typer.Option("--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"),
] = 65536,
cache_setup: Annotated[
bool,
@@ -660,8 +617,7 @@ def train(
Optional[bool],
typer.Option(
"--wandb/--no-wandb",
help="Log per-epoch training metrics to Weights & Biases "
"(requires `uv sync --extra wandb`)",
help="Log per-epoch training metrics to Weights & Biases (requires `uv sync --extra wandb`)",
),
] = None,
wandb_project: Annotated[
@@ -692,8 +648,7 @@ def train(
batch_size_value = int(batch_size)
except ValueError:
typer.echo(
f"error: --batch-size must be an integer or 'auto', "
f"got {batch_size!r}",
f"error: --batch-size must be an integer or 'auto', got {batch_size!r}",
err=True,
)
raise typer.Exit(1)
@@ -753,9 +708,7 @@ def train(
"decoder": stage2_decoder.value if stage2_decoder is not None else None,
"k_max": stage2_k_max,
"context_dim": stage2_context_dim,
"stage1_context": stage2_stage1_context.value
if stage2_stage1_context is not None
else None,
"stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
}.items()
if v is not None
}
@@ -834,9 +787,7 @@ def train(
):
stage_wgan = {**shared_wgan_overrides, **stage_specific}
if stage_wgan:
overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(
stage_wgan
)
overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(stage_wgan)
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
gconfig.validate_config(cfg)
@@ -845,19 +796,13 @@ def train(
_device = torch.device(device) if device else gconfig.auto_device()
if batch_size_auto:
est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
cfg, training=True, stage="stage1"
)
est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(cfg, training=True, stage="stage1")
try:
t["batch_size"] = gconfig.estimate_batch_size(
est_hidden_dim, est_n_blocks, _device
)
t["batch_size"] = gconfig.estimate_batch_size(est_hidden_dim, est_n_blocks, _device)
except ValueError as exc:
typer.echo(f"error: {exc}", err=True)
raise typer.Exit(1)
typer.echo(
f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
)
typer.echo(f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)")
if out is not None:
out_dir = out
@@ -913,39 +858,19 @@ def new_run(
n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
stage1_generator: Annotated[
Optional[Mode], typer.Option("--stage1-generator")
] = None,
stage1_hidden_dim: Annotated[
Optional[int], typer.Option("--stage1-hidden-dim")
] = None,
stage1_n_res_blocks: Annotated[
Optional[int], typer.Option("--stage1-n-res-blocks")
] = None,
stage1_generator: Annotated[Optional[Mode], typer.Option("--stage1-generator")] = None,
stage1_hidden_dim: Annotated[Optional[int], typer.Option("--stage1-hidden-dim")] = None,
stage1_n_res_blocks: Annotated[Optional[int], typer.Option("--stage1-n-res-blocks")] = None,
stage1_dropout: Annotated[Optional[float], typer.Option("--stage1-dropout")] = None,
stage2_generator: Annotated[
Optional[Mode], typer.Option("--stage2-generator")
] = None,
stage2_hidden_dim: Annotated[
Optional[int], typer.Option("--stage2-hidden-dim")
] = None,
stage2_n_res_blocks: Annotated[
Optional[int], typer.Option("--stage2-n-res-blocks")
] = None,
stage2_generator: Annotated[Optional[Mode], typer.Option("--stage2-generator")] = None,
stage2_hidden_dim: Annotated[Optional[int], typer.Option("--stage2-hidden-dim")] = None,
stage2_n_res_blocks: Annotated[Optional[int], typer.Option("--stage2-n-res-blocks")] = None,
stage2_dropout: Annotated[Optional[float], typer.Option("--stage2-dropout")] = None,
stage2_decoder: Annotated[
Optional[Decoder], typer.Option("--stage2-decoder")
] = None,
stage2_decoder: Annotated[Optional[Decoder], typer.Option("--stage2-decoder")] = None,
stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None,
stage2_context_dim: Annotated[
Optional[int], typer.Option("--stage2-context-dim")
] = None,
stage2_stage1_context: Annotated[
Optional[Stage1Context], typer.Option("--stage2-stage1-context")
] = None,
conditioning: Annotated[
Optional[Conditioning], typer.Option("--conditioning")
] = None,
stage2_context_dim: Annotated[Optional[int], typer.Option("--stage2-context-dim")] = None,
stage2_stage1_context: Annotated[Optional[Stage1Context], typer.Option("--stage2-stage1-context")] = None,
conditioning: Annotated[Optional[Conditioning], typer.Option("--conditioning")] = None,
router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None,
@@ -956,16 +881,13 @@ def new_run(
] = None,
comment: Annotated[
Optional[str],
typer.Option(
"--comment", help="Free-text note recorded in config.toml's meta section"
),
typer.Option("--comment", help="Free-text note recorded in config.toml's meta section"),
] = None,
data: Annotated[
Optional[Path],
typer.Option(
"--data",
help="Dataset path to fill in the printed next-step command "
"(not stored in the config)",
help="Dataset path to fill in the printed next-step command (not stored in the config)",
),
] = None,
force: Annotated[
@@ -977,9 +899,7 @@ def new_run(
] = False,
dry_run: Annotated[
bool,
typer.Option(
"--dry-run", help="Print the resolved config without writing anything"
),
typer.Option("--dry-run", help="Print the resolved config without writing anything"),
] = False,
) -> None:
"""Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir.
@@ -1030,9 +950,7 @@ def new_run(
"decoder": stage2_decoder.value if stage2_decoder is not None else None,
"k_max": stage2_k_max,
"context_dim": stage2_context_dim,
"stage1_context": stage2_stage1_context.value
if stage2_stage1_context is not None
else None,
"stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
}.items()
if v is not None
}
@@ -1116,9 +1034,7 @@ def new_run(
@app.command()
def predict(
data: Annotated[
Path, typer.Argument(help="Parquet file or directory of parquet files")
],
data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
checkpoint: Annotated[
Path,
typer.Option(
@@ -1151,8 +1067,7 @@ def predict(
typer.Option(
"--batch-size",
"-b",
help="Inference batch size, or 'auto' to estimate from free GPU "
"memory (cuda devices only)",
help="Inference batch size, or 'auto' to estimate from free GPU memory (cuda devices only)",
),
] = "4096",
steps: Annotated[
@@ -1221,8 +1136,7 @@ def predict(
if "sec_phys" not in ckpt.get("normalizer", {}):
typer.echo(
"error: checkpoint has no normalizer.sec_phys — retrain with the "
"current code",
"error: checkpoint has no normalizer.sec_phys — retrain with the current code",
err=True,
)
raise typer.Exit(1)
@@ -1235,9 +1149,7 @@ def predict(
stage2_k_max = _stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
if batch_size_auto:
est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
model_cfg, training=False
)
est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg, training=False)
try:
batch_size_value = gconfig.estimate_batch_size(
est_hidden_dim,
@@ -1248,9 +1160,7 @@ def predict(
except ValueError as exc:
typer.echo(f"error: {exc}", err=True)
raise typer.Exit(1)
typer.echo(
f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
)
typer.echo(f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)")
assert batch_size_value is not None
bs = batch_size_value
@@ -1311,16 +1221,10 @@ def predict(
return iter_file_chunks(path, offset=offset, k_max=stage2_k_max)
return iter_cond_chunks(path, offset=offset)
cond_pdg_topn = (
pdg_topn_map.class_map if particle_conditioning == "onehot" else None
)
cond_mat_topn = (
mat_topn_map.class_map if material_conditioning == "onehot" else None
)
cond_pdg_topn = pdg_topn_map.class_map if particle_conditioning == "onehot" else None
cond_mat_topn = mat_topn_map.class_map if material_conditioning == "onehot" else None
def _concat(
a: dict[str, np.ndarray], b: dict[str, np.ndarray]
) -> dict[str, np.ndarray]:
def _concat(a: dict[str, np.ndarray], b: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
return {k: np.concatenate([a[k], b[k]], axis=0) for k in a}
def _process(piece: dict[str, np.ndarray]) -> None:
@@ -1352,19 +1256,13 @@ def predict(
cc = torch.from_numpy(cond_cont).float().to(_device)
ck = torch.from_numpy(cond_cat).long().to(_device)
stage1_norm, n_sec_pred = sample_stage1(
model, cc, ck, steps=steps, ddpm_steps=stage1_ddpm_steps
)
stage1_norm, n_sec_pred = sample_stage1(model, cc, ck, steps=steps, ddpm_steps=stage1_ddpm_steps)
if coord == Coord.global_:
# A fresh v0.3.0 Stage1Model owns no n_sec_head —
# sample_stage1 returns n_sec_pred=None then, so ask stage 2.
n_sec_pred = resolve_n_sec(
model, sec_decoder, cc, ck, stage1_norm, n_sec_pred
)
sec_cont, sec_type, _sec_valid_pred = sample_stage2(
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
)
n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred)
sec_cont, sec_type, _sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
n_sec_pred_np = n_sec_pred.cpu().numpy()
pred = stage1_norm.cpu().numpy() # normalised
@@ -1387,14 +1285,8 @@ def predict(
"material": piece["material"],
"layer_id": piece["layer_id"],
"n_sec": piece["n_sec"],
**{
f"pred_{name}": raw[:, j]
for j, name in enumerate(LOCAL_TARGET_NAMES)
},
**{
f"true_{name}": target_raw[:, j]
for j, name in enumerate(LOCAL_TARGET_NAMES)
},
**{f"pred_{name}": raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
**{f"true_{name}": target_raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
}
)
else:
@@ -1404,9 +1296,7 @@ def predict(
# (hence delta_e == edep + e_sec) holds by construction. e_sec_pred
# doubles as the stick-breaking energy budget for the Stage-2 decode
# below, since the model has no other source for it at inference.
edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(
raw[:, 1:3], piece["pre_E"]
)
edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(raw[:, 1:3], piece["pre_E"])
# Normalise predicted direction then rotate back to world frame
post_dir_local = raw[:, 3:6].copy()
@@ -1419,43 +1309,31 @@ def predict(
travel_dir_local = raw[:, 6:9].copy()
norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True)
travel_dir_local /= np.where(norms < 1e-8, 1.0, norms)
post_pos_world = reconstruct_post_pos(
piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
)
post_pos_world = reconstruct_post_pos(piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local)
# particle_type.target="physical": sec_pdg_code is a reporting-
# only nearest-known-PDG label (never fed back into the model —
# "no snapping at inference"). "onehot"/"embedding": PDG
# resolution IS the secondary's identity — see
# decode_secondary_identity's docstring.
sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, _l1_dist = (
decode_secondary_identity(
sec_decoder,
sec_cont,
sec_type,
n_sec_pred_np,
e_sec_pred,
piece["pre_dir"],
sec_phys_norm,
pdg_map,
pdg_topn_map,
other_policy,
None,
)
sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, _l1_dist = decode_secondary_identity(
sec_decoder,
sec_cont,
sec_type,
n_sec_pred_np,
e_sec_pred,
piece["pre_dir"],
sec_phys_norm,
pdg_map,
pdg_topn_map,
other_policy,
None,
)
sec_pdg_list = [
sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_pdg_list = [sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
sec_dx_list = [
sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_dy_list = [
sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_dz_list = [
sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_dx_list = [sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)]
sec_dy_list = [sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)]
sec_dz_list = [sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)]
table = pa.table(
{
@@ -1537,9 +1415,7 @@ def predict(
typer.echo(f"reference: {ref_path}")
if skipped:
codes = ", ".join(
f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items())
)
codes = ", ".join(f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items()))
typer.echo(
f"warning: skipped {skipped:,} row(s) with unknown PDG code(s): {codes}",
err=True,
@@ -1587,9 +1463,7 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
@app.command()
def rollout(
data: Annotated[
Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")
],
data: Annotated[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)"),
@@ -1609,9 +1483,7 @@ def rollout(
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,
max_steps: Annotated[int, typer.Option("--max-steps", help="Max steps per individual track")] = 1000,
steps: Annotated[
int,
typer.Option(
@@ -1629,9 +1501,7 @@ def rollout(
"requires a checkpoint trained with EMA enabled.",
),
] = Weights.raw,
batch_size: Annotated[
int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")
] = 4096,
batch_size: Annotated[int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")] = 4096,
max_tracks_per_event: Annotated[
Optional[int],
typer.Option(
@@ -1646,15 +1516,9 @@ def rollout(
help="Override the oracle's NN-distance escape threshold [mm]",
),
] = None,
n_events: Annotated[
Optional[int], typer.Option("--n-events", help="Cap number of seed events")
] = None,
device: Annotated[
Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")
] = None,
out: Annotated[
Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")
] = None,
n_events: Annotated[Optional[int], typer.Option("--n-events", help="Cap number of seed events")] = None,
device: Annotated[Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")] = None,
out: Annotated[Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")] = None,
seed: Annotated[
Optional[int],
typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
@@ -1679,8 +1543,7 @@ def rollout(
if "sec_phys" not in ckpt.get("normalizer", {}):
typer.echo(
"error: checkpoint has no normalizer.sec_phys — retrain with the "
"current code",
"error: checkpoint has no normalizer.sec_phys — retrain with the current code",
err=True,
)
raise typer.Exit(1)
@@ -1730,10 +1593,7 @@ def rollout(
typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
oracle = GeometryOracle.load(geometry)
typer.echo(
f"loaded geometry oracle: {geometry} "
f"(escape_threshold={oracle.escape_threshold:.3f})"
)
typer.echo(f"loaded geometry oracle: {geometry} (escape_threshold={oracle.escape_threshold:.3f})")
files = find_parquet_files(data)
seeds = _seed_from_data(files, n_events)
@@ -1821,9 +1681,7 @@ def rollout(
# stage2_model.particle_type.target="embedding"; omitted (not
# written as null) otherwise, so giant.analysis can tell "not
# applicable to this checkpoint" apart from "collector empty".
**(
{"type_embedding_l1_dist": l1_summary} if l1_summary is not None else {}
),
**({"type_embedding_l1_dist": l1_summary} if l1_summary is not None else {}),
# Full architecture spec baked into the checkpoint — includes the
# entire router sub-dict, not just a hand-picked subset, so any
# model knob (router type/n_experts, noise_dim, vocab sizes, ...)
@@ -1855,9 +1713,7 @@ app.add_typer(analyze_app, name="analyze")
def analyze_prep(
rollout_yaml: Annotated[
Path,
typer.Argument(
help="giant rollout YAML sidecar (names the rollout + reference files)"
),
typer.Argument(help="giant rollout YAML sidecar (names the rollout + reference files)"),
],
run_dir: Annotated[
Optional[Path],
@@ -1872,9 +1728,7 @@ def analyze_prep(
top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
chunks: Annotated[
int,
typer.Option(
"--chunks", help="Split each plot's data into this many event_id chunks"
),
typer.Option("--chunks", help="Split each plot's data into this many event_id chunks"),
] = 1,
) -> None:
"""Read the rollout YAML → shared.json + run_meta.json in the run directory."""
@@ -1894,15 +1748,9 @@ def analyze_prep(
@analyze_app.command("compute-one")
def analyze_compute_one(
id: Annotated[
str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
],
run_dir: Annotated[
Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
],
chunk: Annotated[
int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")
] = 0,
id: Annotated[str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")],
run_dir: Annotated[Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")],
chunk: Annotated[int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")] = 0,
) -> None:
"""Run one (plot, chunk)'s streaming reduction (this is what each condor job runs)."""
from giant.analysis import compute_one
@@ -1913,12 +1761,8 @@ def analyze_compute_one(
@analyze_app.command("merge-one")
def analyze_merge_one(
id: Annotated[
str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
],
run_dir: Annotated[
Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
],
id: Annotated[str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")],
run_dir: Annotated[Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")],
) -> None:
"""Merge one plot's chunk partials into its final reduced JSON.
@@ -1942,14 +1786,10 @@ def analyze_list() -> None:
@analyze_app.command("render")
def analyze_render(
run_dir: Annotated[
Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")
],
run_dir: Annotated[Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")],
gallery: Annotated[
bool,
typer.Option(
"--gallery/--no-gallery", help="Run `gallery generate` after rendering"
),
typer.Option("--gallery/--no-gallery", help="Run `gallery generate` after rendering"),
] = False,
) -> None:
"""Render reduced artifacts to styled PDFs + gallery metadata (local; needs LaTeX)."""
@@ -1971,9 +1811,7 @@ def analyze_submit(
help="Override the run directory (default: <cwd>/analysis_runs/analysis_<id>)",
),
] = None,
docker_image: Annotated[
str, typer.Option("--docker-image")
] = "cverstege/alma9-gridjob",
docker_image: Annotated[str, typer.Option("--docker-image")] = "cverstege/alma9-gridjob",
request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 8192,
remote: Annotated[
bool,
@@ -1989,9 +1827,7 @@ def analyze_submit(
n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
dry_run: Annotated[
bool, typer.Option("--dry-run", help="Write files but don't condor_submit")
] = False,
dry_run: Annotated[bool, typer.Option("--dry-run", help="Write files but don't condor_submit")] = False,
) -> None:
"""prep + write the HTCondor submit description (one job per plot x chunk), then submit."""
import subprocess
giant/config.py
+6 -25
View File
@@ -299,13 +299,7 @@ DEFAULT_CONFIG: dict = {
def git_hash() -> str:
try:
return (
subprocess.check_output(
["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
)
.decode()
.strip()
)
return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
except Exception:
return "unknown"
@@ -357,12 +351,8 @@ def estimate_batch_size(
calibration since there's no backward graph or optimizer state.
"""
if device.type != "cuda":
raise ValueError(
f"--batch-size auto is only supported on cuda devices, got {device.type!r}"
)
device_index = (
device.index if device.index is not None else torch.cuda.current_device()
)
raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}")
device_index = device.index if device.index is not None else torch.cuda.current_device()
free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
if training:
ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
@@ -690,11 +680,7 @@ def validate_config(cfg: dict) -> None:
for stage_name in ("stage1_model", "stage2_model"):
router = _get_path(cfg, f"{stage_name}.router") or {}
if (
router.get("enabled")
and router.get("type") in ("pdg", "process")
and particle_type == "physical"
):
if router.get("enabled") and router.get("type") in ("pdg", "process") and particle_type == "physical":
raise ValueError(
f"{stage_name}.router.type = {router['type']!r} builds its "
"own training-vocab-scoped embedding, incompatible with "
@@ -704,9 +690,7 @@ def validate_config(cfg: dict) -> None:
"conditioning.particle.type"
)
if _get_path(cfg, "stage2_model.router.tie_to_stage1") and not _get_path(
cfg, "stage1_model.active"
):
if _get_path(cfg, "stage2_model.router.tie_to_stage1") and not _get_path(cfg, "stage1_model.active"):
raise ValueError(
"stage2_model.router.tie_to_stage1 = true requires "
"stage1_model.active = true (there is no stage-1 router to tie to)"
@@ -737,10 +721,7 @@ def validate_config(cfg: dict) -> None:
if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
history = _get_path(cfg, "stage2_model.autoregressive.history")
if history not in ("markov", "attention"):
raise ValueError(
f"stage2_model.autoregressive.history = {history!r} — must "
"be 'markov' or 'attention'"
)
raise ValueError(f"stage2_model.autoregressive.history = {history!r} — must be 'markov' or 'attention'")
teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
if teacher_forcing not in ("always", "scheduled", "never"):
raise ValueError(
giant/data/dataset.py
+1 -3
View File
@@ -123,9 +123,7 @@ class StreamingStepsDataset(IterableDataset):
buf_n = 0
for path in files:
for chunk in iter_file_chunks(
path, offset=self._offsets[path], k_max=self.k_max
):
for chunk in iter_file_chunks(path, offset=self._offsets[path], k_max=self.k_max):
mask = sorted_membership(chunk["event_id"], self._events_arr)
if not mask.any():
continue
giant/data/loader.py
+8 -24
View File
@@ -115,9 +115,7 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
return out
def _df_to_dict(
df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX
) -> dict[str, np.ndarray]:
def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
has_sec_lists = "sec_E_list" in df.columns
d: dict[str, np.ndarray] = {
@@ -136,9 +134,7 @@ def _df_to_dict(
# / 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)
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),
@@ -151,16 +147,12 @@ def _df_to_dict(
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
)
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, offset: int = 0, k_max: int = K_MAX
) -> dict[str, np.ndarray]:
def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
@@ -170,9 +162,7 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
return _offset_event_id(ids, offset)
def iter_file_chunks(
path: str | Path, offset: int = 0, k_max: int = K_MAX
) -> Iterator[dict[str, np.ndarray]]:
def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> Iterator[dict[str, np.ndarray]]:
"""Yield one parquet row-group at a time so a large file never fully loads.
`k_max` sets the padded width of the sec_*_list columns (should match
@@ -214,15 +204,11 @@ def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]
}
def iter_cond_chunks(
path: str | Path, offset: int = 0
) -> Iterator[dict[str, np.ndarray]]:
def iter_cond_chunks(path: str | Path, offset: int = 0) -> 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(), offset=offset
)
yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)
def build_index_maps(
@@ -315,9 +301,7 @@ class TopNMap:
other_members: dict
def build_topn_map_from_files(
files: list[Path], column: str, n_classes: int, cast=str
) -> TopNMap:
def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
"""Scan `column` and build a frequency-capped value->index map, structurally
identical to `build_process_map_from_files` (shares its ranking core via
`_topn_plus_other_map`), generalized over the source column and key type.
giant/data/setup_cache.py
+12 -39
View File
@@ -108,10 +108,7 @@ def normalizer_key(
# affect which cond_cont columns are computed for real vs. zero-filled
# (giant.data.transforms._physical_cond_columns), so both must be part of
# the key or two mixed-axis runs could collide on the same cache entry.
return (
f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}"
f"_mcond={material_conditioning}"
)
return f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}_mcond={material_conditioning}"
# Top-N-map axes: "pdg" keys match pdg_map's int
@@ -126,10 +123,7 @@ def topn_key(axis: str, n_classes: int) -> str:
sidecar stays reusable across runs with different emb_dim (see the
dict[int, dict] precedent `proc_maps` sets, keyed by n_experts)."""
if axis not in _TOPN_AXIS_CASTS:
raise ValueError(
f"unknown top-N map axis {axis!r}, expected one of "
f"{sorted(_TOPN_AXIS_CASTS)}"
)
raise ValueError(f"unknown top-N map axis {axis!r}, expected one of {sorted(_TOPN_AXIS_CASTS)}")
return f"{axis}:{n_classes}"
@@ -165,9 +159,7 @@ class NormalizerEntry:
"tgt_norm": self.tgt_norm.to_dict(),
"sec_phys_norm": self.sec_phys_norm.to_dict(),
"n_train_steps": self.n_train_steps,
"energy_quantiles": np.asarray(
self.energy_quantiles, dtype=np.float32
).tolist(),
"energy_quantiles": np.asarray(self.energy_quantiles, dtype=np.float32).tolist(),
}
@classmethod
@@ -234,13 +226,8 @@ class SetupCache:
np.array(d["event_index"]["counts"], dtype=np.int64),
)
proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
normalizers = {
k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
}
topn_maps = {
k: topnmap_from_json(v, axis=k.split(":", 1)[0])
for k, v in d.get("topn_maps", {}).items()
}
normalizers = {k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()}
topn_maps = {k: topnmap_from_json(v, axis=k.split(":", 1)[0]) for k, v in d.get("topn_maps", {}).items()}
return cls(
fingerprint=d["fingerprint"],
git_hash=d.get("git_hash", "unknown"),
@@ -263,18 +250,14 @@ class SetupCache:
fingerprint=other.fingerprint,
git_hash=other.git_hash,
vocab=other.vocab if other.vocab is not None else self.vocab,
event_index=(
other.event_index if other.event_index is not None else self.event_index
),
event_index=(other.event_index if other.event_index is not None else self.event_index),
proc_maps={**self.proc_maps, **other.proc_maps},
normalizers={**self.normalizers, **other.normalizers},
topn_maps={**self.topn_maps, **other.topn_maps},
)
def load(
data: str | Path, files: list[Path], echo=lambda *a, **k: None
) -> SetupCache | None:
def load(data: str | Path, files: list[Path], echo=lambda *a, **k: None) -> SetupCache | None:
"""Load and validate the sidecar for `data`; `None` on any miss (never raises).
A missing file, corrupt JSON, format-version mismatch, dimension-constant
@@ -298,9 +281,7 @@ def load(
echo("setup cache: format version changed — ignoring stale cache")
return None
if raw.get("dims") != _DIMS:
echo(
"setup cache: model dimension constants changed — ignoring stale cache"
)
echo("setup cache: model dimension constants changed — ignoring stale cache")
return None
fp = fingerprint_files(files)
if raw.get("fingerprint") != fp:
@@ -343,9 +324,7 @@ def save(
with open(lock_path, "a") as lock_file:
fcntl.flock(lock_file, fcntl.LOCK_EX)
try:
base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(
files
)
base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files)
merged = base.merge(sections)
payload = json.dumps(merged.to_json(), separators=(",", ":"))
tmp.write_text(payload)
@@ -353,9 +332,7 @@ def save(
finally:
fcntl.flock(lock_file, fcntl.LOCK_UN)
except OSError as exc:
echo(
f"setup cache: could not write {path} ({exc}) — continuing without caching"
)
echo(f"setup cache: could not write {path} ({exc}) — continuing without caching")
try:
tmp.unlink(missing_ok=True)
except OSError:
@@ -366,16 +343,12 @@ def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.nd
"""Unique event ids + per-event row (step) counts, across all `files`."""
if not files:
return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
all_ids = np.concatenate(
[load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
)
all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
unique_ids, counts = np.unique(all_ids, return_counts=True)
return unique_ids, counts
def n_train_steps_for_split(
unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
) -> int:
def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:
"""Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.
`train_events_arr` must be ascending and duplicate-free (as produced by
giant/data/transforms.py
+21 -63
View File
@@ -84,9 +84,7 @@ def energy_simplex_encode(
return z.astype(np.float32)
def energy_simplex_decode(
z: np.ndarray, pre_E: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
def energy_simplex_decode(z: np.ndarray, pre_E: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.
A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
@@ -118,9 +116,7 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
arbitrary second operand) on every row; profiling on a 114M-row file
showed `np.cross` as the single hottest call inside this rotation.
"""
axis = np.stack(
[pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
)
axis = np.stack([pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1)
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
# axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
# pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
@@ -199,9 +195,7 @@ def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarra
kxv = _cross_with_z_axis(axis, post_dir) # (N,3)
kdv = (axis * post_dir).sum(axis=1, keepdims=True) # (N,1)
return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
np.float32
)
return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
class Normalizer:
@@ -341,9 +335,7 @@ def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
return sorted_arr[idx] == values
def _vectorized_map_lookup(
values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0
) -> np.ndarray:
def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0) -> np.ndarray:
"""Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.
Replaces a per-element Python dict lookup with one `searchsorted` call.
@@ -387,9 +379,7 @@ def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
disp = post_pos - pre_pos
norm = np.linalg.norm(disp, axis=1, keepdims=True)
safe_norm = np.where(norm < 1e-7, 1.0, norm)
return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
np.float32
)
return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(np.float32)
def reconstruct_post_pos(
@@ -408,9 +398,7 @@ def reconstruct_post_pos(
return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)
def inv_local_frame_rotation(
pre_dir: np.ndarray, post_dir_local: np.ndarray
) -> np.ndarray:
def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> np.ndarray:
"""Inverse of local_frame_rotation: rotate from local frame back to world frame.
Applies R^T (same axis, negative angle) to post_dir_local.
@@ -424,9 +412,7 @@ def inv_local_frame_rotation(
kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)
# Negative angle: sin_t → -sin_t
return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
np.float32
)
return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
_STICK_LOGIT_CLIP = 10.0 # logit value used for the last valid secondary slot
@@ -491,14 +477,10 @@ def encode_secondaries(
remaining_raw = e_sec - cumsum[:, i - 1]
shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
remaining = np.maximum(remaining_raw, _EPS)
f = np.clip(
sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS
)
f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
logit = np.log(f / (1.0 - f)).astype(np.float32)
# Last valid slot: give it the full remaining budget
is_last = sec_valid[:, i] & ~(
sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
)
is_last = sec_valid[:, i] & ~(sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool))
logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
logit = np.where(
sec_valid[:, i],
@@ -525,9 +507,7 @@ def encode_secondaries(
# Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
valid_mask = sec_valid[:, i]
if valid_mask.any():
dir_local[valid_mask, i] = local_frame_rotation(
pre_dir[valid_mask], sec_dir_list[valid_mask, i]
)
dir_local[valid_mask, i] = local_frame_rotation(pre_dir[valid_mask], sec_dir_list[valid_mask, i])
if sec_pdg_list is not None:
from giant.particles import particle_phys_array
@@ -553,9 +533,7 @@ def encode_secondaries(
return sec_cont.astype(np.float32)
def encode_secondary_type_idx(
sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict
) -> np.ndarray:
def encode_secondary_type_idx(sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict) -> np.ndarray:
"""Per-secondary-slot class index into `class_map` — (N, K_MAX) int64.
`class_map` is either a top-N-plus-other map's `class_map`
@@ -583,9 +561,7 @@ def encode_secondary_type_idx(
# isn't guaranteed to be a key — an arbitrary present one always is).
dummy = next(iter(class_map))
safe_pdg = np.where(sec_valid, sec_pdg_list, dummy)
idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(
N, K
)
idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(N, K)
return np.where(sec_valid, idx, 0).astype(np.int64)
@@ -659,9 +635,7 @@ def decode_secondary_cont(
for i in range(K):
valid = sec_valid[:, i]
if valid.any():
sec_dir_world[valid, i] = inv_local_frame_rotation(
pre_dir[valid], dir_local[valid, i]
)
sec_dir_world[valid, i] = inv_local_frame_rotation(pre_dir[valid], dir_local[valid, i])
return sec_E, sec_dir_world, sec_valid
@@ -704,9 +678,7 @@ def decode_secondaries(
sec_cont = sec_cont.copy()
sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
sec_cont, n_sec, e_sec, pre_dir
)
sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont, n_sec, e_sec, pre_dir)
log_mass = sec_cont[:, :, 4] # (N, K)
charge = sec_cont[:, :, 5] # (N, K)
@@ -760,16 +732,12 @@ def _physical_cond_columns(
elif particle_conditioning in ("embedding", "onehot"):
particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
else:
raise ValueError(
f"unknown conditioning.particle.type {particle_conditioning!r}"
)
raise ValueError(f"unknown conditioning.particle.type {particle_conditioning!r}")
if material_conditioning == "physical":
from giant.materials import material_properties_array
z_eff, a_eff, density, x0, lambda_int = material_properties_array(
data["material"]
).T
z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T
material_cols = np.column_stack(
[
z_eff,
@@ -782,9 +750,7 @@ def _physical_cond_columns(
elif material_conditioning in ("embedding", "onehot"):
material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
else:
raise ValueError(
f"unknown conditioning.material.type {material_conditioning!r}"
)
raise ValueError(f"unknown conditioning.material.type {material_conditioning!r}")
return np.column_stack([particle_cols, material_cols]).astype(np.float32)
@@ -848,9 +814,7 @@ def build_cond_features(
cond_cat = np.column_stack(cat_cols)
if cond_normalizer is not None:
cond_cont = _cond_normalizer_transform(
cond_cont, cond_normalizer, particle_conditioning, material_conditioning
)
cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, particle_conditioning, material_conditioning)
return cond_cont, cond_cat
@@ -972,13 +936,9 @@ def build_features(
"""
post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
travel_dir_local = local_frame_rotation(
data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
)
travel_dir_local = local_frame_rotation(data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"]))
energy_z = energy_simplex_encode(
data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
) # (N, 2)
energy_z = energy_simplex_encode(data["edep"], data["e_sec"], data["post_E"], data["pre_E"]) # (N, 2)
target_s1 = np.column_stack(
[
@@ -1014,9 +974,7 @@ def build_features(
cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
cond_cat = np.column_stack(cat_cols) # (N, 2/3/4)
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
# Secondary continuous targets
sec_E_list = data.get("sec_E_list")
giant/geometry.py
+7 -29
View File
@@ -31,10 +31,7 @@ import numpy as np
import pandas as pd
import pyarrow.parquet as pq
_INSTALL_HINT = (
"the geometry oracle needs scikit-learn — install it with "
"`uv sync --extra cpu --extra geometry`"
)
_INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"
def _require_sklearn():
@@ -63,9 +60,7 @@ class _SlabLookup:
layer_ids: np.ndarray # (n_segments,) int64, layer_id of each segment
radius_max: float # largest transverse radius seen in training data
def query(
self, pos: np.ndarray, margin: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
def query(self, pos: np.ndarray, margin: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
other = [i for i in range(3) if i != self.axis]
z = pos[:, self.axis]
radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)
@@ -75,11 +70,7 @@ class _SlabLookup:
material = self.materials[idx]
layer_id = self.layer_ids[idx]
escaped = (
(z < self.z_edges[0] - margin)
| (z > self.z_edges[-1] + margin)
| (radius > self.radius_max + margin)
)
escaped = (z < self.z_edges[0] - margin) | (z > self.z_edges[-1] + margin) | (radius > self.radius_max + margin)
return material, layer_id, escaped
@@ -300,16 +291,12 @@ def _fit_slab_lookup(
"""
other = [i for i in range(3) if i != axis]
z = pos[:, axis].astype(np.float64)
radius = np.sqrt(
pos[:, other[0]].astype(np.float64) ** 2
+ pos[:, other[1]].astype(np.float64) ** 2
)
radius = np.sqrt(pos[:, other[0]].astype(np.float64) ** 2 + pos[:, other[1]].astype(np.float64) ** 2)
z_min, z_max = float(z.min()), float(z.max())
if z_min == z_max:
raise ValueError(
"all points share the same depth-axis coordinate — pick a "
"different `depth_axis` or use method='knn'/'svm'"
"all points share the same depth-axis coordinate — pick a different `depth_axis` or use method='knn'/'svm'"
)
edges = np.linspace(z_min, z_max, n_bins + 1)
bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
@@ -344,12 +331,7 @@ def _fit_slab_lookup(
bin_layer = bin_layer[fill_from]
# Run-length-encode consecutive bins sharing a label into segments.
changed = (
np.flatnonzero(
(bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
)
+ 1
)
changed = np.flatnonzero((bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])) + 1
seg_starts = np.concatenate([[0], changed])
z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
materials = bin_material[seg_starts]
@@ -460,11 +442,7 @@ def build_geometry_oracle(
# Escape threshold from the reference point spacing. Sample a subset for the
# median 2-NN distance (the 1st neighbour of a training point is itself).
nn = NearestNeighbors(n_neighbors=2).fit(X)
probe = (
X
if len(X) <= 20_000
else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
)
probe = X if len(X) <= 20_000 else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
d2, _ = nn.kneighbors(probe, n_neighbors=2)
median_nn = float(np.median(d2[:, 1]))
escape_threshold = escape_factor * median_nn
giant/materials.py
+9 -26
View File
@@ -64,18 +64,10 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
"G4_CESIUM_IODIDE": MaterialProperties(
z_eff=54.0, a_eff=129.904539, density=4.51, x0=1.860288, lambda_int=39.305990
),
"G4_Pb": MaterialProperties(
z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950
),
"G4_W": MaterialProperties(
z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580
),
"G4_Cu": MaterialProperties(
z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940
),
"G4_Fe": MaterialProperties(
z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300
),
"G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950),
"G4_W": MaterialProperties(z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580),
"G4_Cu": MaterialProperties(z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940),
"G4_Fe": MaterialProperties(z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300),
"G4_BRASS": MaterialProperties(
z_eff=30.939130,
a_eff=68.500857,
@@ -83,9 +75,7 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
x0=1.367465,
lambda_int=16.947420,
),
"G4_POLYSTYRENE": MaterialProperties(
z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880
),
"G4_POLYSTYRENE": MaterialProperties(z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880),
"G4_PLASTIC_SC_VINYLTOLUENE": MaterialProperties(
z_eff=3.368421,
a_eff=6.219791,
@@ -108,20 +98,15 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
x0=30392.070000,
lambda_int=71009.500000,
),
"G4_lAr": MaterialProperties(
z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400
),
"G4_lAr": MaterialProperties(z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400),
}
def get_material_properties(
name: str, table: dict[str, MaterialProperties] | None = None
) -> MaterialProperties:
def get_material_properties(name: str, table: dict[str, MaterialProperties] | None = None) -> MaterialProperties:
t = MATERIAL_PROPERTIES if table is None else table
if name not in t:
raise UnknownMaterialError(
f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES "
f"-- add it (known: {sorted(t)})"
f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES -- add it (known: {sorted(t)})"
)
props = t[name]
if any(v is None for v in props):
@@ -134,9 +119,7 @@ def get_material_properties(
return props
def material_properties_array(
names: np.ndarray, table: dict[str, MaterialProperties] | None = None
) -> np.ndarray:
def material_properties_array(names: np.ndarray, table: dict[str, MaterialProperties] | None = None) -> np.ndarray:
"""(N,) str material names -> (N, 5) float32 [z_eff, a_eff, density, x0, lambda_int]."""
out = np.array(
[get_material_properties(str(m), table) for m in np.asarray(names)],
giant/model/network.py
+58 -194
View File
@@ -30,11 +30,7 @@ class SinusoidalEmbedding(nn.Module):
super().__init__()
assert dim % 2 == 0, "dim must be even"
half = dim // 2
freqs = torch.exp(
-math.log(10000)
* torch.arange(half, dtype=torch.float32)
/ max(half - 1, 1)
)
freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
self.register_buffer("freqs", freqs)
def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -43,9 +39,7 @@ class SinusoidalEmbedding(nn.Module):
return torch.cat([args.sin(), args.cos()], dim=-1) # (B, dim)
def cat_col_layout(
particle_type: str, material_type: str
) -> tuple[int | None, int | None]:
def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]:
"""`cond_cat` column indices for each axis's top-N-onehot index, or
`None` if that axis isn't `"onehot"`.
@@ -119,18 +113,14 @@ class ConditionEncoder(nn.Module):
super().__init__()
self.particle_cfg = dict(particle_cfg)
self.material_cfg = dict(material_cfg)
self._particle_topn_col, self._material_topn_col = cat_col_layout(
particle_cfg["type"], material_cfg["type"]
)
self._particle_topn_col, self._material_topn_col = cat_col_layout(particle_cfg["type"], material_cfg["type"])
p_type = particle_cfg["type"]
p_emb_dim = particle_cfg["emb_dim"]
if p_type == "embedding":
self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim)
elif p_type == "physical":
self.particle_mlp = _make_axis_mlp(
PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1)
)
self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1))
elif p_type != "onehot":
raise ValueError(f"unknown conditioning.particle.type {p_type!r}")
@@ -139,9 +129,7 @@ class ConditionEncoder(nn.Module):
if m_type == "embedding":
self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim)
elif m_type == "physical":
self.material_mlp = _make_axis_mlp(
MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1)
)
self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1))
elif m_type != "onehot":
raise ValueError(f"unknown conditioning.material.type {m_type!r}")
@@ -157,9 +145,7 @@ class ConditionEncoder(nn.Module):
if p_type == "embedding":
return self.pdg_emb(cond_cat[:, 0])
if p_type == "physical":
particle_phys = cond_cont[
:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
]
particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
return self.particle_mlp(particle_phys)
assert self._particle_topn_col is not None
return F.one_hot(
@@ -245,9 +231,7 @@ class Router(nn.Module):
"""(B, n_experts) soft weights, rows summing to 1."""
raise NotImplementedError
def combine_weights(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
"""(B, n_experts) train-time expert-combination weights.
Default (`gumbel=False`): identical to `gate()`. Opt-in
@@ -265,16 +249,12 @@ class Router(nn.Module):
"""(B,) hard expert index, used for eval-time grouped dispatch."""
return self.gate(cond_cont, cond_cat).argmax(dim=-1)
def balance_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
"""Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,)
return (importance.std() / (importance.mean() + 1e-8)) ** 2
def classify_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
"""Optional supervised auxiliary loss shaping the router's own belief.
Default: none (a scalar 0). Routers gating on an unobservable
@@ -282,16 +262,12 @@ class Router(nn.Module):
"""
return torch.zeros((), device=cond_cont.device)
def entropy_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
"""Optional auxiliary loss rewarding sharper (lower-entropy) routing."""
norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
return norm_entropy
def gate_stats(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for
the full explanation, unchanged in v0.3.0."""
gate = self.gate(cond_cont, cond_cat) # (B, n_experts)
@@ -321,9 +297,7 @@ def build_router(name: str, n_experts: int, **kwargs) -> Router:
in one config without special-casing.
"""
if name not in ROUTER_REGISTRY:
raise ValueError(
f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
)
raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
cls = ROUTER_REGISTRY[name]
accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
filtered = {k: v for k, v in kwargs.items() if k in accepted}
@@ -374,8 +348,7 @@ class EnergyRouter(Router):
if learn_width or learn_temperature:
if not (width_min_ratio < 1.0 < width_max_ratio):
raise ValueError(
f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
f"({width_max_ratio}) must bracket 1.0"
f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
)
self._width_lo = width_min_ratio * temperature
self._width_hi = width_max_ratio * temperature
@@ -388,10 +361,7 @@ class EnergyRouter(Router):
centers = torch.linspace(-2.0, 2.0, n_experts)
else:
if len(centers_init) != n_experts:
raise ValueError(
f"centers_init has {len(centers_init)} values, "
f"expected n_experts={n_experts}"
)
raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
centers = torch.tensor(list(centers_init), dtype=torch.float32)
if learn_centers:
self.centers = nn.Parameter(centers)
@@ -436,9 +406,7 @@ class PdgRouter(Router):
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim)
d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
-1
) # (B, n_experts)
d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1) # (B, n_experts)
return torch.softmax(-d2 / self.temperature, dim=-1)
@@ -477,9 +445,7 @@ class ProcessRouter(Router):
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
def classify_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
@@ -500,14 +466,10 @@ class ComposedRouter(Router):
joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0)
for router in self.routers[1:]:
g = router.gate(cond_cont, cond_cat) # (B, n_i)
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
1
) # (B, prod so far)
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1) # (B, prod so far)
return joint
def classify_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
total = torch.zeros((), device=cond_cont.device)
for router in self.routers:
total = total + router.classify_loss(cond_cont, cond_cat, labels)
@@ -560,9 +522,7 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]:
_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
def _check_router_conditioning_compat(
router_types: list[str], particle_conditioning: str
) -> None:
def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None:
"""Reject a router axis that reintroduces a training-vocab PDG lookup
under `conditioning.particle.type = "physical"`.
@@ -596,16 +556,12 @@ def _build_router_from_cfg(
shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
if router_cfg["type"] == "composed":
axes = _parse_composed_axes(router_cfg)
_check_router_conditioning_compat(
[a["type"] for a in axes], particle_conditioning
)
_check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning)
router = build_composed_router(axes, **shared_vocab)
router.gumbel = bool(router_cfg.get("gumbel", False))
return router
_check_router_conditioning_compat([router_cfg["type"]], particle_conditioning)
router_kwargs = {
k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
}
router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
router_kwargs.setdefault("pdg_vocab", pdg_vocab)
router_kwargs.setdefault("mat_vocab", mat_vocab)
router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
@@ -637,9 +593,7 @@ class ExpertTrunk(nn.Module):
) -> None:
super().__init__()
self.input_proj = nn.Linear(in_dim, hidden_dim)
self.blocks = nn.ModuleList(
[ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_proj = nn.Linear(hidden_dim, out_dim)
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
@@ -710,12 +664,7 @@ class MonolithicTrunk(Trunk):
) -> None:
super().__init__()
self.input_proj = nn.Linear(in_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, cond_dim, dropout=dropout)
for _ in range(n_res_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_res_blocks)])
self.out_proj = nn.Linear(hidden_dim, out_dim)
def forward(
@@ -745,12 +694,7 @@ class RoutedTrunk(Trunk):
super().__init__()
self.router = router
self.experts = nn.ModuleList(
[
ExpertTrunk(
in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout
)
for _ in range(router.n_experts)
]
[ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) for _ in range(router.n_experts)]
)
def forward(
@@ -760,9 +704,7 @@ class RoutedTrunk(Trunk):
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
) -> torch.Tensor:
return _route_forward(
self.experts, self.router, x, cond, cond_cont, cond_cat, self.training
)
return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training)
def build_trunk(
@@ -775,9 +717,7 @@ def build_trunk(
dropout: float = 0.0,
) -> Trunk:
if router is not None:
return RoutedTrunk(
router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout
)
return RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
@@ -845,13 +785,9 @@ class _CausalAttnBlock(nn.Module):
def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None:
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(
dim, n_heads, dropout=dropout, batch_first=True
)
self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
self.norm2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
)
self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))
def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
h = self.norm1(x)
@@ -860,9 +796,7 @@ class _CausalAttnBlock(nn.Module):
x = x + self.mlp(self.norm2(x))
return x
def step(
self, x_new: torch.Tensor, kv_cache: torch.Tensor | None
) -> tuple[torch.Tensor, torch.Tensor]:
def step(self, x_new: torch.Tensor, kv_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
"""`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None`
(first position) or `(B, T, dim)` — `norm1(x)` of every earlier
position at this same block. Returns `(out, new_kv_cache)`, `out`
@@ -900,15 +834,11 @@ class AttentionHistory(HistoryEncoder):
`Stage2Autoregressive.history_step`.
"""
def __init__(
self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2
) -> None:
def __init__(self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2) -> None:
super().__init__()
self.start = nn.Parameter(torch.zeros(in_dim))
self.in_proj = nn.Linear(in_dim, out_dim)
self.blocks = nn.ModuleList(
[_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)]
)
self.blocks = nn.ModuleList([_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)])
def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
start = self.start.view(1, 1, -1).expand_as(feat)
@@ -959,9 +889,7 @@ def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
def stage2_trunk_sec_dim(
particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int
) -> int:
def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int:
"""`Stage2OneShot`'s trunk output width.
`target = "physical"` is untouched from v0.2/today:
@@ -1020,17 +948,13 @@ class Stage1Model(nn.Module):
self.cond_enc = (
cond_enc
if cond_enc is not None
else ConditionEncoder(
pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
)
else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
)
has_time = generator in ("flow", "ddpm")
self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
in_dim = noise_dim if generator == "wgan" else x_dim
self.trunk = build_trunk(
router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
)
self.trunk = build_trunk(router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
self.n_sec_head = None
if n_sec_head_k_max is not None:
self.n_sec_head = nn.Sequential(
@@ -1047,16 +971,10 @@ class Stage1Model(nn.Module):
t: torch.Tensor | None = None,
) -> torch.Tensor:
c_emb = self.cond_enc(cond_cont, cond_cat)
cond = (
torch.cat([self.time_emb(t), c_emb], dim=-1)
if self.time_emb is not None
else c_emb
)
cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
return self.trunk(x_t, cond, cond_cont, cond_cat)
def predict_n_sec(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
"""Return n_sec logits (B, K_MAX+1) from conditioning alone. Only
valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
configs predict n_sec from Stage2OneShot instead."""
@@ -1127,9 +1045,7 @@ class Stage2OneShot(nn.Module):
self.cond_enc = (
cond_enc
if cond_enc is not None
else ConditionEncoder(
pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
)
else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
)
self.context_adapter = ContextAdapter(x_dim, context_dim)
self.fuse = nn.Sequential(
@@ -1140,9 +1056,7 @@ class Stage2OneShot(nn.Module):
self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
in_dim = noise_dim if generator == "wgan" else sec_dim
self.trunk = build_trunk(
router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
)
self.trunk = build_trunk(router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
self.n_sec_head = None
if build_n_sec_head:
self.n_sec_head = nn.Sequential(
@@ -1162,9 +1076,7 @@ class Stage2OneShot(nn.Module):
self._type_k_max = k_max
self._type_emb_dim = emb_dim
def _cond_embed(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
) -> torch.Tensor:
def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
base = self.cond_enc(cond_cont, cond_cat)
ctx = self.context_adapter(stage1_out)
return self.fuse(torch.cat([base, ctx], dim=-1))
@@ -1178,11 +1090,7 @@ class Stage2OneShot(nn.Module):
t: torch.Tensor | None = None,
) -> torch.Tensor:
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
cond = (
torch.cat([self.time_emb(t), c_emb], dim=-1)
if self.time_emb is not None
else c_emb
)
cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
return self.trunk(x_t, cond, cond_cont, cond_cat)
def predict_n_sec(
@@ -1274,10 +1182,7 @@ class Stage2Autoregressive(nn.Module):
) -> None:
super().__init__()
if history not in ("markov", "attention"):
raise ValueError(
f"stage2_model.autoregressive.history={history!r} — must be "
"'markov' or 'attention'"
)
raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'")
self.history_kind = history
self.generator_kind = generator
self.noise_dim = noise_dim
@@ -1289,9 +1194,7 @@ class Stage2Autoregressive(nn.Module):
self.cond_enc = (
cond_enc
if cond_enc is not None
else ConditionEncoder(
pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
)
else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
)
self.context_adapter = ContextAdapter(x_dim, context_dim)
self.base_fuse = nn.Sequential(
@@ -1305,15 +1208,11 @@ class Stage2Autoregressive(nn.Module):
history_dim = cond_out_dim
hist_in_dim = CONT_SLOT_DIM + self.type_dim
self.history_encoder: HistoryEncoder = (
AttentionHistory(
hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
)
AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers)
if history == "attention"
else MarkovHistory(hist_in_dim, history_dim)
)
token_fuse_in = (
cond_out_dim + context_dim + history_dim + 2
) # +2: remaining_frac, slot_idx
token_fuse_in = cond_out_dim + context_dim + history_dim + 2 # +2: remaining_frac, slot_idx
self.token_fuse = nn.Sequential(
nn.Linear(token_fuse_in, cond_out_dim),
nn.SiLU(),
@@ -1350,9 +1249,7 @@ class Stage2Autoregressive(nn.Module):
nn.Linear(hidden_dim // 2, self.type_dim),
)
def _base_cond(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
) -> torch.Tensor:
def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
base = self.cond_enc(cond_cont, cond_cat)
ctx = self.context_adapter(stage1_out)
return self.base_fuse(torch.cat([base, ctx], dim=-1))
@@ -1392,9 +1289,7 @@ class Stage2Autoregressive(nn.Module):
return self.history_encoder.init_cache()
return None
def history_step(
self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache
) -> tuple[torch.Tensor, object]:
def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
"""One inference slot's worth of history encoding: advances `cache`
(from `init_history_cache`, or a previous `history_step` call) by
`token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
@@ -1445,9 +1340,7 @@ class Stage2Autoregressive(nn.Module):
out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
return out.view(B, K, -1)
def predict_n_sec(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
) -> torch.Tensor:
def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
if self.n_sec_head is None:
raise RuntimeError(
"this Stage2Autoregressive has no n_sec_head — it belongs to "
@@ -1514,9 +1407,7 @@ class CriticModel(nn.Module):
if stage not in ("stage1", "stage2"):
raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
self.stage = stage
self.cond_enc = ConditionEncoder(
pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
)
self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
if stage == "stage2":
self.context_adapter = ContextAdapter(context_in_dim, context_dim)
self.fuse = nn.Sequential(
@@ -1524,12 +1415,7 @@ class CriticModel(nn.Module):
nn.SiLU(),
)
self.input_proj = nn.Linear(in_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
for _ in range(n_res_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)])
self.out_norm = nn.LayerNorm(hidden_dim)
self.out_proj = nn.Linear(hidden_dim, 1)
@@ -1641,9 +1527,7 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
}
def migrate_legacy_state_dict(
old_stage1_sd: dict, old_stage2_sd: dict
) -> tuple[dict, dict]:
def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple[dict, dict]:
"""Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`,
`SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new
`(Stage1Model, Stage2OneShot)` module structure produced by
@@ -1702,11 +1586,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
common representation. `false` (default) keeps v0.2 behaviour:
independent instances with identical config but independent weights.
"""
cfg = (
model_config
if "stage1_model" in model_config
else _migrate_legacy_model_config(model_config)
)
cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
pdg_vocab = cfg["pdg_vocab"]
mat_vocab = cfg["mat_vocab"]
conditioning = cfg["conditioning"]
@@ -1718,9 +1598,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
cond_out_dim = conditioning.get("out_dim", 128)
shared_cond_enc: ConditionEncoder | None = None
if conditioning.get("share_stages"):
shared_cond_enc = ConditionEncoder(
pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
)
shared_cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
@@ -1728,15 +1606,11 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
if s1cfg.get("active", True):
router_cfg = s1cfg.get("router") or {}
if router_cfg.get("enabled"):
stage1_router = _build_router_from_cfg(
router_cfg, pdg_vocab, mat_vocab, particle_conditioning
)
stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
generator = s1cfg.get("generator", "flow")
gen_sub = s1cfg.get(generator, {}) or {}
legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
n_sec_head_k_max = (
s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None
)
n_sec_head_k_max = s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None
result["stage1"] = Stage1Model(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
@@ -1762,9 +1636,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
if router_cfg.get("tie_to_stage1") and stage1_router is not None:
stage2_router = stage1_router
else:
stage2_router = _build_router_from_cfg(
router_cfg, pdg_vocab, mat_vocab, particle_conditioning
)
stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
generator = s2cfg.get("generator", "wgan")
gen_sub = s2cfg.get(generator, {}) or {}
legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
@@ -1796,9 +1668,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
cond_enc=shared_cond_enc,
)
else:
sec_dim = stage2_trunk_sec_dim(
particle_type_cfg, generator, k_max, particle_cfg["emb_dim"]
)
sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, k_max, particle_cfg["emb_dim"])
result["stage2"] = Stage2OneShot(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
@@ -1828,11 +1698,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
"wgan"` training. Training-only — never persisted for inference the way
`build_models`'s pair is. `None` for a stage that's inactive or not
WGAN."""
cfg = (
model_config
if "stage1_model" in model_config
else _migrate_legacy_model_config(model_config)
)
cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
pdg_vocab = cfg["pdg_vocab"]
mat_vocab = cfg["mat_vocab"]
conditioning = cfg["conditioning"]
@@ -1861,9 +1727,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
if s2cfg.get("active", True) and s2cfg.get("generator") == "wgan":
k_max = s2cfg.get("k_max", K_MAX)
particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"}
in_dim = stage2_trunk_sec_dim(
particle_type_cfg, "wgan", k_max, particle_cfg["emb_dim"]
)
in_dim = stage2_trunk_sec_dim(particle_type_cfg, "wgan", k_max, particle_cfg["emb_dim"])
result["stage2"] = CriticModel(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
giant/model/schedule.py
+1 -3
View File
@@ -11,9 +11,7 @@ class CosineSchedule:
steps = np.arange(T + 1, dtype=np.float64)
f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
alpha_bars = (f / f[0]).astype(np.float32)
betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(
np.float32
)
betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)
self.betas = torch.from_numpy(betas)
self.alphas = torch.from_numpy((1.0 - betas))
giant/particles.py
+3 -10
View File
@@ -58,9 +58,7 @@ def particle_mass_charge(pdg: int) -> tuple[float, float]:
if _pdgid.is_nucleus(pdg):
z, a = _pdgid.Z(pdg), _pdgid.A(pdg)
if z is None or a is None:
raise ValueError(
f"PDG {pdg}: is_nucleus but Z/A decode failed"
) from None
raise ValueError(f"PDG {pdg}: is_nucleus but Z/A decode failed") from None
return float(a) * _AMU_MEV, float(z)
raise ValueError(
f"PDG code {pdg} could not be resolved via the `particle` package "
@@ -109,9 +107,7 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd
if len(resolved) == 0:
raise ValueError("nearest_known_pdg: no resolvable candidates")
codes = np.array([r[0] for r in resolved], dtype=np.int64)
table_log_mass = np.log(
np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS
)
table_log_mass = np.log(np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS)
table_charge = np.array([r[2] for r in resolved], dtype=np.float64)
mass = np.asarray(mass, dtype=np.float64)
@@ -180,10 +176,7 @@ def decode_topn_class(
n_other = int(other_mask.sum())
if n_other:
if not topn_map.other_members:
raise ValueError(
"decode_topn_class: 'other' class predicted but "
"topn_map.other_members is empty"
)
raise ValueError("decode_topn_class: 'other' class predicted but topn_map.other_members is empty")
members = np.array(list(topn_map.other_members.keys()), dtype=np.int64)
counts = np.array(list(topn_map.other_members.values()), dtype=np.float64)
if other_policy == "drop":
giant/pipeline.py
+27 -73
View File
@@ -73,21 +73,14 @@ def _seed_energy_router(
if not active:
return
if energy_quantiles.size == 0:
echo(
" warning: no energy samples collected — EnergyRouter falls back to "
"default centers"
)
echo(" warning: no energy samples collected — EnergyRouter falls back to default centers")
return
assert cond_norm.mean is not None and cond_norm.std is not None
levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
energy_idx
]
centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
echo(
f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
)
echo(f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}")
def run_setup_stage(
@@ -142,23 +135,14 @@ def run_setup_stage(
if cache is not None:
cache.event_index = (unique_ids, counts)
train_events, val_events = make_event_split(
unique_ids, val_fraction=val_fraction, seed=seed
)
train_events, val_events = make_event_split(unique_ids, val_fraction=val_fraction, seed=seed)
events_arr = np.array(sorted(train_events))
n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
echo(
f" {int(counts.sum()):,} steps | "
f"{len(train_events)} train events | "
f"{len(val_events)} val events"
)
echo(f" {int(counts.sum()):,} steps | {len(train_events)} train events | {len(val_events)} val events")
if cache is not None and cache.vocab is not None:
pdg_map, mat_map = cache.vocab
echo(
f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
f"{len(mat_map)} materials)"
)
echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
else:
echo("building vocabulary maps …")
pdg_map, mat_map = build_index_maps_from_files(files)
@@ -174,11 +158,7 @@ def run_setup_stage(
# need that generality.
proc_map: dict[str, int] | None = None
process_router_cfg = next(
(
r
for r in (stage1_router, stage2_router)
if r.get("enabled") and r.get("type") == "process"
),
(r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
None,
)
if process_router_cfg is not None:
@@ -186,10 +166,7 @@ def run_setup_stage(
cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
if cached_proc_map is not None:
proc_map = cached_proc_map
echo(
f"process vocabulary: cache hit ({len(proc_map)} labels, "
f"{n_experts} experts)"
)
echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
else:
echo("building process vocabulary …")
proc_map = build_process_map_from_files(files, n_experts=n_experts)
@@ -214,16 +191,11 @@ def run_setup_stage(
cached = cache.topn_maps.get(cache_key) if cache is not None else None
if cached is not None:
pdg_topn_map = cached
echo(
f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, "
f"{n_classes} classes)"
)
echo(f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, {n_classes} classes)")
else:
echo("building pdg top-N map …")
pdg_topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
echo(
f" {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes"
)
echo(f" {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes")
if cache is not None:
cache.topn_maps[cache_key] = pdg_topn_map
@@ -234,29 +206,17 @@ def run_setup_stage(
cached = cache.topn_maps.get(cache_key) if cache is not None else None
if cached is not None:
mat_topn_map = cached
echo(
f"material top-N map: cache hit ({len(mat_topn_map.class_map)} "
f"materials, {n_classes} classes)"
)
echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
else:
echo("building material top-N map …")
mat_topn_map = build_topn_map_from_files(
files, "material", n_classes=n_classes, cast=str
)
echo(
f" {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes"
)
mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
echo(f" {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
if cache is not None:
cache.topn_maps[cache_key] = mat_topn_map
energy_router_active = any(
r.get("enabled") and r.get("type") == "energy"
for r in (stage1_router, stage2_router)
)
energy_router_active = any(r.get("enabled") and r.get("type") == "energy" for r in (stage1_router, stage2_router))
energy_idx = 3
norm_key = setup_cache.normalizer_key(
val_fraction, seed, particle_conditioning, material_conditioning
)
norm_key = setup_cache.normalizer_key(val_fraction, seed, particle_conditioning, material_conditioning)
entry = cache.normalizers.get(norm_key) if cache is not None else None
if entry is not None:
@@ -281,27 +241,23 @@ def run_setup_stage(
# router against this same (val_fraction, seed, conditioning) key
# never needs to rescan just to seed centers.
collect_energy_sample = energy_router_active or cache is not None
energy_sampler = (
_ReservoirSampler(capacity=100_000) if collect_energy_sample else None
)
energy_sampler = _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
for i, path in enumerate(files):
for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
mask = sorted_membership(chunk["event_id"], events_arr)
if not mask.any():
continue
chunk_tr = {k: v[mask] for k, v in chunk.items()}
cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = (
build_features(
chunk_tr,
pdg_map,
mat_map,
proc_map=proc_map,
require_secondaries=True,
particle_conditioning=particle_conditioning,
material_conditioning=material_conditioning,
sec_phys_only=True,
k_max=k_max,
)
cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = build_features(
chunk_tr,
pdg_map,
mat_map,
proc_map=proc_map,
require_secondaries=True,
particle_conditioning=particle_conditioning,
material_conditioning=material_conditioning,
sec_phys_only=True,
k_max=k_max,
)
cond_acc.update(cond_cont)
tgt_acc.update(target_s1)
@@ -423,9 +379,7 @@ def run_train_job(
# The secondary type-index map depends on stage2_model.particle_type.target,
# independently of conditioning's own onehot/embedding choice above
# (physical stays untouched/None).
particle_type_target = (
cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
)
particle_type_target = cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
if particle_type_target == "onehot":
assert setup.pdg_topn_map is not None
sec_type_class_map = setup.pdg_topn_map.class_map
giant/rollout.py
+35 -83
View File
@@ -120,9 +120,7 @@ def decode_secondary_identity(
pdg_topn_map: "TopNMap | None",
other_policy: str,
rng: np.random.Generator | None,
) -> tuple[
np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None
]:
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
"""Decode Stage 2's raw (sec_cont, sec_type) output into physical
secondary attributes, branching on `sec_decoder.particle_type_cfg`:
@@ -150,14 +148,12 @@ def decode_secondary_identity(
sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
sec_full, n_sec_np, e_sec, pre_dir, sec_phys_normalizer=sec_phys_norm
)
sec_pdg = nearest_known_pdg(
sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
).reshape(sec_mass.shape)
sec_pdg = nearest_known_pdg(sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()).reshape(
sec_mass.shape
)
return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, None
sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir
)
sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir)
sec_type_np = sec_type.cpu().numpy()
l1_dist = None
@@ -182,12 +178,8 @@ def decode_secondary_identity(
l1_dist = np.where(sec_valid, l1_dist, 0.0).astype(np.float32)
sec_mass, sec_charge = particle_phys_array(sec_pdg.reshape(-1)).T
sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(
np.float32
)
sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(
np.float32
)
sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(np.float32)
sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(np.float32)
sec_pdg = np.where(sec_valid, sec_pdg, 0).astype(np.int64)
return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, l1_dist
@@ -299,13 +291,9 @@ class _Recorder:
RAM until the very end.
"""
def __init__(
self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
) -> None:
def __init__(self, sink: Callable[[dict[str, np.ndarray]], None] | None = None) -> None:
self._sink = sink
self._cols: dict[str, list] | None = (
None if sink is not None else {k: [] for k in _RECORD_KEYS}
)
self._cols: dict[str, list] | None = None if sink is not None else {k: [] for k in _RECORD_KEYS}
self.n_rows = 0
self.termination_reason_counts: Counter[str] = Counter()
@@ -313,10 +301,7 @@ class _Recorder:
n = len(cols["event_id"])
if n == 0:
return
row = {
k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
for k in _RECORD_KEYS
}
row = {k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) for k in _RECORD_KEYS}
self.n_rows += n
reasons = row["termination_reason"]
nonempty = reasons[reasons != ""]
@@ -333,8 +318,7 @@ class _Recorder:
def to_dict(self) -> dict[str, np.ndarray]:
assert self._cols is not None, (
"to_dict() is unavailable when streaming to a sink — use "
"n_rows/termination_reason_counts instead"
"to_dict() is unavailable when streaming to a sink — use n_rows/termination_reason_counts instead"
)
out = {}
for k, chunks in self._cols.items():
@@ -619,33 +603,19 @@ 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
@@ -679,41 +649,27 @@ def _step_chunk(
cond_norm,
particle_conditioning=particle_conditioning,
material_conditioning=material_conditioning,
pdg_topn_map=pdg_topn_map.class_map
if particle_conditioning == "onehot"
else None,
mat_topn_map=mat_topn_map.class_map
if material_conditioning == "onehot"
else None,
pdg_topn_map=pdg_topn_map.class_map if particle_conditioning == "onehot" else None,
mat_topn_map=mat_topn_map.class_map if material_conditioning == "onehot" else None,
)
cc = torch.from_numpy(cond_cont).float().to(device)
ck = torch.from_numpy(cond_cat).long().to(device)
stage1_norm, n_sec_pred_stage1 = sample_stage1(
stage1_model, cc, ck, steps, stage1_ddpm_steps
)
stage1_norm, n_sec_pred_stage1 = sample_stage1(stage1_model, cc, ck, steps, stage1_ddpm_steps)
raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
step_length = inv_log_transform(raw[:, 0])
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_pred = resolve_n_sec(
stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1
)
n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)
n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
# --- Secondaries ---
@@ -724,23 +680,19 @@ def _step_chunk(
# decode_secondary_identity's docstring for how each
# particle_type.target differs on whether PDG resolution is a real
# identity decision or just a reporting label.
sec_cont, sec_type, _valid = sample_stage2(
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
)
sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = (
decode_secondary_identity(
sec_decoder,
sec_cont,
sec_type,
n_sec_np,
e_sec,
tr["pre_dir"],
sec_phys_norm,
pdg_map,
pdg_topn_map,
other_policy,
rng,
)
sec_cont, sec_type, _valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
sec_decoder,
sec_cont,
sec_type,
n_sec_np,
e_sec,
tr["pre_dir"],
sec_phys_norm,
pdg_map,
pdg_topn_map,
other_policy,
rng,
)
sec_valid = np.arange(sec_E.shape[1])[None, :] < n_sec_np[:, None]
if l1_dist_collector is not None and sec_type_l1_dist is not None:
giant/sample.py
+12 -36
View File
@@ -67,9 +67,7 @@ def sample_ddpm(
alpha = schedule.alphas[i]
alpha_bar = schedule.alpha_bars[i]
z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
x = (1.0 / alpha.sqrt()) * (
x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
) + beta.sqrt() * z
x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
@@ -173,9 +171,7 @@ def _decode_stage2_flat(
else:
sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
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, sec_valid
@@ -206,9 +202,7 @@ def sample_secondaries(
v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
x = x + v * dt
return _decode_stage2_flat(
sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
)
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
@torch.no_grad()
@@ -226,9 +220,7 @@ def sample_secondaries_wgan(
B = cond_cont.size(0)
z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
return _decode_stage2_flat(
sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
)
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
@torch.no_grad()
@@ -306,12 +298,8 @@ def sample_secondaries_ar(
has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
history_feat = prev_repr.unsqueeze(1) # (B, 1, CONT_SLOT_DIM + type_dim)
remaining_frac = remaining.unsqueeze(1) # (B, 1)
slot_idx = torch.full(
(B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
)
hist, history_cache = sec_decoder.history_step(
history_feat, has_prev, history_cache
)
slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
if generator == "wgan":
z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
@@ -366,21 +354,15 @@ def sample_secondaries_ar(
sec_type[:, k] = type_k
if target == "onehot":
type_for_history = F.one_hot(
type_k.argmax(dim=-1), num_classes=type_dim
).float()
type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
else:
type_for_history = type_k
stick_fraction = torch.sigmoid(cont_k[:, 0])
prev_repr = torch.cat(
[stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1
)
prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
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, sec_valid
@@ -426,16 +408,10 @@ def sample_stage2(
case, so there's nothing to dispatch to here.
"""
if isinstance(sec_decoder, Stage2Autoregressive):
return sample_secondaries_ar(
sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
)
return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
if sec_decoder.generator_kind == "wgan":
return sample_secondaries_wgan(
sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
)
return sample_secondaries(
sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
)
return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
def resolve_n_sec(
giant/training/loop.py
+14 -46
View File
@@ -78,9 +78,7 @@ def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs)
return validate_marginals(model, val_loader, device=device, **kwargs)
def _marginal_kl(
trainers: dict[str, StageTrainer], val_loader, device, **kwargs
) -> float:
def _marginal_kl(trainers: dict[str, StageTrainer], val_loader, device, **kwargs) -> float:
"""Mean marginal KL over the stage-1 sampling chain, or NaN when stage 1
is inactive or `validate_marginals` declined to produce a result."""
stage1 = trainers.get("stage1")
@@ -90,9 +88,7 @@ def _marginal_kl(
stage1,
val_loader,
device,
sec_decoder=trainers["stage2"].sampling_model()
if "stage2" in trainers
else None,
sec_decoder=trainers["stage2"].sampling_model() if "stage2" in trainers else None,
**kwargs,
)
if result is None:
@@ -137,9 +133,7 @@ def train(
trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches)
if not trainers:
raise ValueError(
"no active stage — stage1_model.active and stage2_model.active are both false"
)
raise ValueError("no active stage — stage1_model.active and stage2_model.active are both false")
has_adversarial = any(not tr.supports_val_loss for tr in trainers.values())
checkpoint_extras = {
@@ -147,12 +141,8 @@ def train(
"pdg_map": pdg_map,
"mat_map": mat_map,
"proc_map": proc_map,
"pdg_topn_map": topnmap_to_json(pdg_topn_map)
if pdg_topn_map is not None
else None,
"mat_topn_map": topnmap_to_json(mat_topn_map)
if mat_topn_map is not None
else None,
"pdg_topn_map": topnmap_to_json(pdg_topn_map) if pdg_topn_map is not None else None,
"mat_topn_map": topnmap_to_json(mat_topn_map) if mat_topn_map is not None else None,
"model_config": model_config,
}
@@ -166,10 +156,7 @@ def train(
best_val_loss = ckpt.get("best_val_loss", float("inf"))
global_step = ckpt.get("global_step", 0)
if start_epoch > epochs:
print(
f"checkpoint already completed epoch {start_epoch - 1} "
f"(>= --epochs {epochs}) — nothing to train"
)
print(f"checkpoint already completed epoch {start_epoch - 1} (>= --epochs {epochs}) — nothing to train")
return
collector = MetricsCollector.create(
@@ -206,10 +193,7 @@ def train(
for batch in bar:
B = batch[0].size(0)
collector.add_train_batch(
{
name: trainer.step(batch, device, global_step)
for name, trainer in trainers.items()
},
{name: trainer.step(batch, device, global_step) for name, trainer in trainers.items()},
B,
)
bar.set_postfix_str(collector.postfix(), refresh=False)
@@ -220,9 +204,7 @@ def train(
bar.close()
if shutdown.requested:
ckpt = build_checkpoint(
trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras
)
ckpt = build_checkpoint(trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras)
torch.save(ckpt, out_dir / "last.pt")
last_completed_epoch = epoch - 1
print(
@@ -244,15 +226,10 @@ def train(
break
B = batch[0].size(0)
collector.add_val_batch(
{
name: tr.val_loss(batch, device)
for name, tr in scored.items()
},
{name: tr.val_loss(batch, device) for name, tr in scored.items()},
B,
)
collector.observe_routers(
batch[0].to(device), batch[1].to(device), B
)
collector.observe_routers(batch[0].to(device), batch[1].to(device), B)
# An adversarial stage has no averageable validation loss, so it
# needs the marginal-KL signal every epoch to pick a best
@@ -264,10 +241,7 @@ def train(
elif validate_every > 0 and epoch % validate_every == 0:
stage1 = trainers.get("stage1")
ddpm_steps = 1000
if (
isinstance(stage1, FlowDDPMStageTrainer)
and stage1.ddpm_schedule is not None
):
if isinstance(stage1, FlowDDPMStageTrainer) and stage1.ddpm_schedule is not None:
ddpm_steps = stage1.ddpm_schedule.T
marginal_kl = _marginal_kl(
trainers,
@@ -292,22 +266,16 @@ def train(
collector.set("val/marginal_kl", marginal_kl)
collector.set(
"gpu_mem_mb",
torch.cuda.max_memory_allocated(device) / (1024 * 1024)
if device.type == "cuda"
else 0.0,
)
collector.set(
"samples_per_sec", collector.train_samples / max(epoch_time, 1e-8)
torch.cuda.max_memory_allocated(device) / (1024 * 1024) if device.type == "cuda" else 0.0,
)
collector.set("samples_per_sec", collector.train_samples / max(epoch_time, 1e-8))
collector.set("is_best", int(is_best))
collector.set("epoch_time_s", epoch_time)
print(collector.summary_line(val_loss, epoch_time, is_best))
collector.write_epoch(global_step)
ckpt = build_checkpoint(
trainers, epoch, global_step, best_val_loss, checkpoint_extras
)
ckpt = build_checkpoint(trainers, epoch, global_step, best_val_loss, checkpoint_extras)
if is_best:
best_val_loss = val_loss
ckpt["best_val_loss"] = best_val_loss
giant/training/metrics.py
+10 -31
View File
@@ -116,11 +116,7 @@ class _RouterAccumulator:
def add(self, entropy: torch.Tensor, importance: torch.Tensor, n: int) -> None:
self.entropy += entropy.item() * n
self.importance = (
importance.clone()
if self.importance is None
else self.importance + importance
)
self.importance = importance.clone() if self.importance is None else self.importance + importance
self.n += n
def stats(self) -> dict[str, float]:
@@ -168,22 +164,18 @@ class MetricsCollector:
self._train = {name: _Accumulator() for name in trainers}
self._val = {name: _Accumulator() for name in trainers}
self._routers = {
name: _RouterAccumulator(tr.router.n_experts)
for name, tr in trainers.items()
if tr.router is not None
name: _RouterAccumulator(tr.router.n_experts) for name, tr in trainers.items() if tr.router is not None
}
# Only "mean" specs need summing; "last" specs are read straight off
# the accumulator's most recent stats dict. "grad_norm" is always
# summed — it feeds the run-level `grad_norm` column whether or not
# a trainer reports it per stage.
self._train_keys = {
name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"}
| {"grad_norm"}
name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"} | {"grad_norm"}
for name, tr in trainers.items()
}
self._val_keys = {
name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"}
for name, tr in trainers.items()
name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"} for name, tr in trainers.items()
}
self._run_values: dict[str, float] = {}
self._epoch = 0
@@ -224,13 +216,9 @@ class MetricsCollector:
import wandb
except ImportError as exc:
raise RuntimeError(
"train.wandb = true (--wandb) requires the 'wandb' package — "
"install it via `uv sync --extra wandb`"
"train.wandb = true (--wandb) requires the 'wandb' package — install it via `uv sync --extra wandb`"
) from exc
param_counts = {
name: sum(p.numel() for p in tr.model.parameters())
for name, tr in trainers.items()
}
param_counts = {name: sum(p.numel() for p in tr.model.parameters()) for name, tr in trainers.items()}
param_counts["total"] = sum(param_counts.values())
wandb_run = wandb.init(
project=wandb_project,
@@ -290,9 +278,7 @@ class MetricsCollector:
self._batch_grad_norm += stage_stats.get("grad_norm", 0.0)
if self._ema_seeded:
self._ema_loss += _EMA_ALPHA * (self._batch_loss - self._ema_loss)
self._ema_grad_norm += _EMA_ALPHA * (
self._batch_grad_norm - self._ema_grad_norm
)
self._ema_grad_norm += _EMA_ALPHA * (self._batch_grad_norm - self._ema_grad_norm)
else:
self._ema_loss = self._batch_loss
self._ema_grad_norm = self._batch_grad_norm
@@ -303,9 +289,7 @@ class MetricsCollector:
self._val[name].add(stage_stats, self._val_keys[name], batch_size)
@torch.no_grad()
def observe_routers(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int
) -> None:
def observe_routers(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int) -> None:
"""Record gate diagnostics for every routed stage on this batch.
Called from the validation pass only (as in v0.2/v0.3.0), so a stage
@@ -339,9 +323,7 @@ class MetricsCollector:
payload[f"batch/{name}/lr"] = trainer.optimizer.param_groups[0]["lr"]
if trainer.router is not None:
with torch.no_grad():
entropy, _ = trainer.router.gate_stats(
batch[0].to(device), batch[1].to(device)
)
entropy, _ = trainer.router.gate_stats(batch[0].to(device), batch[1].to(device))
payload[f"batch/{name}/router/entropy"] = entropy.item()
self.wandb_run.log(payload, step=global_step)
@@ -367,10 +349,7 @@ class MetricsCollector:
self._run_values[column] = value
def summary_line(self, val_loss: float, epoch_time: float, is_best: bool) -> str:
bits = [
trainer.summary(self.train_means(name))
for name, trainer in self.trainers.items()
]
bits = [trainer.summary(self.train_means(name)) for name, trainer in self.trainers.items()]
marker = " [best]" if is_best else ""
return (
f"epoch {self._epoch:{self.epoch_width}d}/{self.epochs} "
giant/training/stage2_inputs.py
+9 -21
View File
@@ -100,9 +100,9 @@ def _assemble_stage2_real(
`Stage2Autoregressive`'s per-token target also uses; the two must stay in
lockstep. See `_assemble_stage2_ar_target`'s docstring for the
(target, generator) width rules."""
return _assemble_stage2_ar_target(
sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
).flatten(1)
return _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim).flatten(
1
)
def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor:
@@ -137,18 +137,14 @@ def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
def _ar_meta(
k_max: int, batch: int, device: torch.device, fraction: torch.Tensor
) -> dict[str, torch.Tensor]:
def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]:
"""`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR
conditioning tensors that don't depend on *which* history representation
(ground truth vs. the scheduled-sampling mix) produced `fraction`.
Shared by `_assemble_stage2_ar_inputs` and
`_assemble_stage2_ar_inputs_scheduled`, which differ only in
`history_feat`."""
slot_idx = (
torch.arange(k_max, device=device).float() / max(k_max - 1, 1)
).unsqueeze(0)
slot_idx = (torch.arange(k_max, device=device).float() / max(k_max - 1, 1)).unsqueeze(0)
return {
"has_prev": _ar_has_prev(k_max, device).expand(batch, -1),
"remaining_frac": _remaining_energy_fraction(fraction),
@@ -182,9 +178,7 @@ def _assemble_stage2_ar_inputs(
return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)}
def _stage2_tf_prob(
mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int
) -> float:
def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
"""P(condition slot k+1 on the TRUE token k rather than the model's own
prediction), for the current epoch
(`stage2_model.autoregressive.teacher_forcing`).
@@ -260,9 +254,7 @@ def _assemble_stage2_ar_inputs_scheduled(
device = sec_cont.device
B, K = sec_cont.shape[0], sec_cont.shape[1]
if p_tf >= 1.0:
return _assemble_stage2_ar_inputs(
sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim
)
return _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim)
was_training = model.training
sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
@@ -273,9 +265,7 @@ def _assemble_stage2_ar_inputs_scheduled(
fraction_gt = _stick_fraction(sec_cont)
dir_gt = sec_cont[..., 1:CONT_SLOT_DIM]
type_repr_gt = _type_repr(
sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim
)
type_repr_gt = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
fraction_pred, dir_pred, type_repr_pred = _history_repr_from_ar_sample(
sec_cont_pred, sec_type_pred, particle_type_cfg
)
@@ -323,8 +313,6 @@ def _relax_onehot_type_slice(
cont, type_logits = x[..., :cont_dim], x[..., cont_dim:]
if grad_probe is not None:
cont.register_hook(lambda g: grad_probe.__setitem__("cont", g.norm().item()))
type_logits.register_hook(
lambda g: grad_probe.__setitem__("type", g.norm().item())
)
type_logits.register_hook(lambda g: grad_probe.__setitem__("type", g.norm().item()))
type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1)
return torch.cat([cont, type_soft], dim=-1).reshape(B, -1)
giant/training/trainers.py
+30 -98
View File
@@ -41,9 +41,7 @@ from giant.training.stage2_inputs import (
@torch.no_grad()
def _update_ema(
ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
) -> None:
def _update_ema(ema_model: torch.nn.Module, model: torch.nn.Module, decay: float) -> None:
for ema_p, p in zip(ema_model.parameters(), model.parameters()):
ema_p.mul_(decay).add_(p, alpha=1 - decay)
@@ -129,9 +127,7 @@ class StageSpec:
type_gumbel_tau_end: float = 0.1
@classmethod
def from_config(
cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int
) -> "StageSpec":
def from_config(cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int) -> "StageSpec":
t = cfg["train"]
stage_cfg = cfg[f"{name}_model"]
router_cfg = stage_cfg.get("router") or {}
@@ -144,8 +140,7 @@ class StageSpec:
decoder=stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot",
lambda_weight=stage_cfg.get("lambda", 1.0),
n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1),
particle_type=cfg["stage2_model"].get("particle_type")
or {"target": "physical"},
particle_type=cfg["stage2_model"].get("particle_type") or {"target": "physical"},
particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"],
# train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
# (giant/config.py), so they read directly; the field defaults
@@ -244,9 +239,7 @@ class StageTrainer:
# --- schedule -------------------------------------------------------
def _init_lr_schedule(
self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int
) -> None:
def _init_lr_schedule(self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int) -> None:
self._lr_lambda = _cosine_warmup_lambda(warmup_steps, total_steps)
self.total_steps = total_steps
self.lr_sched = optim.lr_scheduler.LambdaLR(optimizer, self._lr_lambda)
@@ -270,9 +263,7 @@ class StageTrainer:
"""This stage's fragment of the end-of-epoch console line."""
raise NotImplementedError
def val_objective(
self, train_means: dict, val_means: dict, marginal_kl: float
) -> float:
def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
"""This stage's contribution to the best-checkpoint selection score."""
raise NotImplementedError
@@ -359,9 +350,7 @@ class StageTrainer:
return target.flatten(1) if flatten else target
@staticmethod
def _sec_mask(
n_sec: torch.Tensor, k_max: int, device: torch.device
) -> torch.Tensor:
def _sec_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> torch.Tensor:
"""`(B, K_MAX)` bool prefix mask: slot k is valid iff `k < n_sec`."""
return torch.arange(k_max, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
@@ -397,9 +386,7 @@ class StageTrainer:
return l_nsec, nsec_acc
@staticmethod
def _step_optimizer(
optimizer: optim.Optimizer, loss: torch.Tensor, params: list
) -> float:
def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
"""`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
the pre-clip grad norm. The one place the grad-clip constant lives."""
optimizer.zero_grad()
@@ -450,9 +437,7 @@ class StageTrainer:
class FlowDDPMStageTrainer(StageTrainer):
"""flow or ddpm generator for a single stage."""
def __init__(
self, spec: StageSpec, model: torch.nn.Module, device: torch.device
) -> None:
def __init__(self, spec: StageSpec, model: torch.nn.Module, device: torch.device) -> None:
if spec.is_stage2 and spec.generator not in ("flow",):
raise NotImplementedError(
f"stage2_model.generator={spec.generator!r} is accepted by the "
@@ -466,26 +451,16 @@ class FlowDDPMStageTrainer(StageTrainer):
# NotImplementedError above): "physical" keeps it folded in
# (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding"
# pull it out into model.type_head instead (0 here).
self._flow_type_dim = (
None
if self.particle_type_cfg.get("target", "physical") == "physical"
else 0
)
self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0
self.params = list(self.model.parameters())
self.optimizer = optim.AdamW(
self.params, lr=spec.lr, weight_decay=spec.weight_decay
)
self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
self._init_lr_schedule(
self.optimizer,
warmup_steps=spec.warmup_epochs * spec.steps_per_epoch,
total_steps=max(spec.epochs * spec.steps_per_epoch, 1),
)
self.ddpm_schedule = (
CosineSchedule(T=spec.ddpm_n_steps).to(device)
if spec.generator == "ddpm"
else None
)
self.ddpm_schedule = CosineSchedule(T=spec.ddpm_n_steps).to(device) if spec.generator == "ddpm" else None
self.train_metrics = [
train_metric(key)
@@ -515,9 +490,7 @@ class FlowDDPMStageTrainer(StageTrainer):
]
self.stage_metrics = [stage_metric("lr")]
def _generator_loss(
self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None
):
def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None):
if not self.is_stage2:
if self.generator == "flow":
return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
@@ -584,22 +557,16 @@ class FlowDDPMStageTrainer(StageTrainer):
mask = sec_mask.float()
denom = mask.sum().clamp(min=1)
if self.particle_type_cfg.get("target") == "onehot":
ce = F.cross_entropy(
type_out.transpose(1, 2), sec_type_idx, reduction="none"
)
ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none")
l_type = (ce * mask).sum() / denom
type_acc = (
(type_out.argmax(-1) == sec_type_idx).float() * mask
).sum() / denom
type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom
else: # "embedding"
target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
se = ((type_out - target_vec) ** 2).mean(-1)
l_type = (se * mask).sum() / denom
return l_type, type_acc
def _compute(
self, batch: tuple, device: torch.device, epoch: int | None = None
) -> dict:
def _compute(self, batch: tuple, device: torch.device, epoch: int | None = None) -> dict:
"""`epoch=None` (the `val_loss` path) always uses full teacher
forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation
should stay a stable, non-stochastic ground-truth comparison; only
@@ -619,23 +586,13 @@ class FlowDDPMStageTrainer(StageTrainer):
x1_s2 = None
ar_inputs = None
if self.is_stage2 and self.decoder == "autoregressive":
ar_inputs = self._ar_inputs(
cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
)
x1_s2 = self._sec_target(
sec_cont, sec_type_idx, self.generator, flatten=False
)
ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False)
elif self.is_stage2:
x1_s2 = self._sec_target(
sec_cont, sec_type_idx, self.generator, flatten=True
)
x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True)
l_gen = self._generator_loss(
cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs
)
l_nsec, nsec_acc = self._n_sec_loss(
cond_cont, cond_cat, stage1_ctx, n_sec, device
)
l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs)
l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
l_type, type_acc = self._type_loss(
cond_cont,
@@ -653,11 +610,7 @@ class FlowDDPMStageTrainer(StageTrainer):
l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
total = (
self.spec.lambda_weight * l_gen
+ self.spec.n_sec_lambda * l_nsec
+ self.particle_type_lambda * l_type
)
total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type
if self.spec.lambda_balance > 0:
total = total + self.spec.lambda_balance * l_balance
if self.spec.lambda_proc > 0:
@@ -698,9 +651,7 @@ class FlowDDPMStageTrainer(StageTrainer):
@torch.no_grad()
def val_loss(self, batch: tuple, device: torch.device) -> dict:
return {
key: value.item() for key, value in self._compute(batch, device).items()
}
return {key: value.item() for key, value in self._compute(batch, device).items()}
# --- reporting ------------------------------------------------------
@@ -710,9 +661,7 @@ class FlowDDPMStageTrainer(StageTrainer):
def summary(self, means: dict) -> str:
return f"{self.name}[loss={means.get('loss', 0.0):.3f}]"
def val_objective(
self, train_means: dict, val_means: dict, marginal_kl: float
) -> float:
def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
return val_means.get("loss", 0.0)
@@ -799,15 +748,8 @@ class WGANStageTrainer(StageTrainer):
if self.decoder == "autoregressive":
epoch = global_step // self.spec.steps_per_epoch
ar = self._ar_inputs(
cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
)
real = (
self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(
B, -1
)
* mask
)
ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(B, -1) * mask
z = torch.randn(B, k_max, self.model.noise_dim, device=device)
fake_raw = self.model(
z,
@@ -891,9 +833,7 @@ class WGANStageTrainer(StageTrainer):
# --- generator (+ n_sec) step ---
did_g_step = global_step % self.n_critic == 0
l_nsec, nsec_acc = self._n_sec_loss(
cond_cont, cond_cat, stage1_ctx, n_sec, device
)
l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
# On a non-generator-step batch with no n_sec_head on this stage
# (n_sec now defaults to stage 2), there's nothing for
@@ -902,9 +842,7 @@ class WGANStageTrainer(StageTrainer):
skip_g_step = not did_g_step and self.model.n_sec_head is None
if did_g_step:
g_loss_adv = generator_loss(critic_fn, fake)
g_loss = (
self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
)
g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
else:
g_loss_adv = torch.zeros((), device=device)
g_loss = self.spec.n_sec_lambda * l_nsec
@@ -941,14 +879,9 @@ class WGANStageTrainer(StageTrainer):
return stats["d_loss"] + stats["g_loss"]
def summary(self, means: dict) -> str:
return (
f"{self.name}[d={means.get('d_loss', 0.0):.3f} "
f"g={means.get('g_loss', 0.0):.3f}]"
)
return f"{self.name}[d={means.get('d_loss', 0.0):.3f} g={means.get('g_loss', 0.0):.3f}]"
def val_objective(
self, train_means: dict, val_means: dict, marginal_kl: float
) -> float:
def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
"""No monotone per-batch WGAN loss fit for averaging, so
best-checkpoint selection uses the real marginal-KL signal when
`validate_marginals` produced one, and falls back to this epoch's own
@@ -1002,8 +935,7 @@ def build_stage_trainers(
if spec.generator == "wgan":
critic = critics.get(name)
assert critic is not None, (
f"{name}_model.generator='wgan' requires a critic (see "
"giant.model.network.build_critics)"
f"{name}_model.generator='wgan' requires a critic (see giant.model.network.build_critics)"
)
trainers[name] = WGANStageTrainer(spec, model, critic, device)
else:
giant/validate.py
+15 -57
View File
@@ -8,9 +8,7 @@ from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
_SEC_PHYS_NAMES = ["log_mass", "charge"]
def _histogram_kl(
p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
) -> float:
def _histogram_kl(p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8) -> float:
"""KL(P || Q) between two 1D samples, estimated via a shared histogram."""
lo = min(p_samples.min(), q_samples.min())
hi = max(p_samples.max(), q_samples.max())
@@ -33,9 +31,7 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
return counts / total if total > 0 else counts
def _categorical_kl(
real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8
) -> float:
def _categorical_kl(real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8) -> float:
"""KL(P_real || Q_gen) between two class-index samples over `n_classes`
categories, estimated from bincount fractions. NaN if either side has no
valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
@@ -48,9 +44,7 @@ def _categorical_kl(
return float(np.sum(p * np.log(p / q)))
def _embedding_nearest_class(
vectors: torch.Tensor, emb_weight: torch.Tensor
) -> np.ndarray:
def _embedding_nearest_class(vectors: torch.Tensor, emb_weight: torch.Tensor) -> np.ndarray:
"""Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
(vocab, emb_dim) — same computation as
`giant.particles.decode_embedding_nearest`, but returning the raw class
@@ -109,11 +103,7 @@ def validate_marginals(
sec_decoder.eval()
k_max = sec_decoder.k_max if sec_decoder is not None else 0
target = (
sec_decoder.particle_type_cfg.get("target", "physical")
if sec_decoder is not None
else "physical"
)
target = sec_decoder.particle_type_cfg.get("target", "physical") if sec_decoder is not None else "physical"
all_real, all_gen = [], []
all_n_sec_real, all_n_sec_pred = [], []
@@ -131,9 +121,7 @@ def validate_marginals(
cond_cont = cond_cont.to(device)
cond_cat = cond_cat.to(device)
gen, n_sec_pred = sample_stage1(
stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps
)
gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)
all_real.append(x1.numpy())
all_gen.append(gen.cpu().numpy())
@@ -141,9 +129,7 @@ def validate_marginals(
if sec_decoder is None:
continue
n_sec_pred = resolve_n_sec(
stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
)
n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
n_sec_pred_np = n_sec_pred.cpu().numpy()
n_sec_np = n_sec.numpy()
all_n_sec_real.append(n_sec_np)
@@ -155,9 +141,7 @@ def validate_marginals(
sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
)
gen_frac = 1.0 / (
1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
)
gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
gen_valid = sec_valid_pred.cpu().numpy()
if target == "physical":
@@ -182,17 +166,9 @@ def validate_marginals(
real = np.concatenate(all_real, axis=0)
generated = np.concatenate(all_gen, axis=0)
kl_divergence = np.array(
[
_histogram_kl(real[:, j], generated[:, j], bins=kl_bins)
for j in range(real.shape[1])
]
)
kl_divergence = np.array([_histogram_kl(real[:, j], generated[:, j], bins=kl_bins) for j in range(real.shape[1])])
header = (
f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} "
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
)
header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
print(f"\n{header}")
print("-" * len(header))
for j, name in enumerate(LOCAL_TARGET_NAMES):
@@ -215,10 +191,7 @@ def validate_marginals(
n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
energy_fraction_kl = np.full(k_max, np.nan)
print(
f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} "
f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
)
print(f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}")
n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
print(n_sec_dist_header)
print("-" * len(n_sec_dist_header))
@@ -240,10 +213,7 @@ def validate_marginals(
continue
kl = _histogram_kl(r, g, bins=kl_bins)
energy_fraction_kl[j] = kl
print(
f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} "
f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}"
)
print(f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}")
result.update(
{
@@ -259,12 +229,7 @@ def validate_marginals(
phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2)
if len(phys_real) > 0 and len(phys_gen) > 0:
phys_kl = np.array(
[
_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
for j in range(2)
]
)
phys_kl = np.array([_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)])
else:
phys_kl = np.full(2, np.nan)
@@ -278,21 +243,14 @@ def validate_marginals(
if len(r) == 0 or len(g) == 0:
continue
print(
f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
)
result.update(
{"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl}
)
result.update({"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl})
else:
type_class_real = np.concatenate(all_type_class_real, axis=0)
type_class_gen = np.concatenate(all_type_class_gen, axis=0)
n_classes = (
sec_decoder.type_dim
if target == "onehot"
else sec_decoder.cond_enc.pdg_emb.weight.size(0)
)
n_classes = sec_decoder.type_dim if target == "onehot" else sec_decoder.cond_enc.pdg_emb.weight.size(0)
type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)
print(
pyproject.toml
+3
View File
@@ -52,6 +52,9 @@ analysis = [
giant = "giant.cli:app"
dwarf = "scripts.dwarf:app"
[tool.ruff]
line-length = 120
[tool.coverage.run]
source = ["giant", "scripts"]
omit = ["*/legacy/*"]
scripts/bump_dataset_version.py
+15 -49
View File
@@ -82,9 +82,7 @@ def _max_index(parent: Path, pattern: re.Pattern) -> int:
def _git_user_name() -> str | None:
try:
out = subprocess.run(
["git", "config", "user.name"], capture_output=True, text=True, timeout=2
)
out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2)
except (OSError, subprocess.SubprocessError):
# OSError (e.g. git not on PATH) and subprocess.SubprocessError
# (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
@@ -142,9 +140,7 @@ def plan_bump_schema(
raw_gen_dir = root / "raw" / kind / gen_tag
processed_gen_dir = root / "processed" / kind / gen_tag
if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
raise SystemExit(
f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
)
raise SystemExit(f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first")
if target is not None:
if not SCHEMA_RE.match(target):
raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
@@ -154,9 +150,7 @@ def plan_bump_schema(
schema_tag = f"schema{next_schema}"
new_dirs = [processed_gen_dir / schema_tag]
by_suffix = f" ({by})" if by else ""
log_line = (
f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
)
log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
return new_dirs, log_line
@@ -212,9 +206,7 @@ def _manifest_referenced_files(pools_root: Path) -> set[Path]:
return referenced
def _referenced_root_count(
raw_gen_dir: Path, processed_gen_dir: Path
) -> tuple[int, int]:
def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
"""(total .root files, count with a same-named .parquet under any schema) for one gen."""
if not raw_gen_dir.is_dir():
return 0, 0
@@ -243,9 +235,7 @@ def _referenced_root_count(
return total, referenced
def _referenced_parquet_count(
schema_dir: Path, manifest_referenced: set[Path]
) -> tuple[int, int]:
def _referenced_parquet_count(schema_dir: Path, manifest_referenced: set[Path]) -> tuple[int, int]:
"""(total .parquet files, count listed in at least one manifest) for one schema dir."""
if not schema_dir.is_dir():
return 0, 0
@@ -342,11 +332,7 @@ def print_status(root: Path) -> None:
grand_files = 0
for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
kind = kind_dir.name
gens = sorted(
int(m.group(1))
for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
if m
)
gens = sorted(int(m.group(1)) for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir()) if m)
print(_colorize(f"{kind}/", "kind"))
kind_total = 0
kind_files = 0
@@ -359,25 +345,18 @@ def print_status(root: Path) -> None:
schemas = sorted(
int(m.group(1))
for m in (
SCHEMA_RE.match(p.name)
for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
if p.is_dir()
SCHEMA_RE.match(p.name) for p in (schema_dir.iterdir() if schema_dir.is_dir() else []) if p.is_dir()
)
if m
)
schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
schema_counts = {
s: _referenced_parquet_count(
schema_dir / f"schema{s}", manifest_referenced
)
for s in schemas
s: _referenced_parquet_count(schema_dir / f"schema{s}", manifest_referenced) for s in schemas
}
processed_size = sum(schema_sizes.values())
processed_files = sum(c[0] for c in schema_counts.values())
processed_referenced = sum(c[1] for c in schema_counts.values())
raw_files, raw_referenced = _referenced_root_count(
raw_gen_dir, processed_gen_dir
)
raw_files, raw_referenced = _referenced_root_count(raw_gen_dir, processed_gen_dir)
gen_total = raw_size + processed_size
gen_files = raw_files + processed_files
kind_total += gen_total
@@ -425,9 +404,7 @@ def print_status(root: Path) -> None:
print(_reason_line(schema_reason, indent=4))
else:
print(_colorize(" (none)", "schema"))
print(
_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files)
)
print(_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files))
print()
grand_total += kind_total
grand_files += kind_files
@@ -526,13 +503,8 @@ def plan_update_manifest(
return result, missing
def apply_update_manifest(
manifest_path: Path, lines: list[tuple[str, str | None]]
) -> None:
out = [
replacement if replacement is not None else original
for original, replacement in lines
]
def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
out = [replacement if replacement is not None else original for original, replacement in lines]
manifest_path.write_text("\n".join(out) + "\n")
@@ -552,9 +524,7 @@ def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
return files
def plan_create_manifest(
output_path: Path, parquet_files: list[Path]
) -> tuple[list[str], list[Path], list[Path]]:
def plan_create_manifest(output_path: Path, parquet_files: list[Path]) -> tuple[list[str], list[Path], list[Path]]:
"""Return (relative_lines, missing_files, resolved_abs_paths)."""
manifest_dir = output_path.resolve().parent
lines: list[str] = []
@@ -569,9 +539,7 @@ def plan_create_manifest(
return lines, missing, resolved
def check_holdout_overlap(
output_path: Path, resolved_new_files: list[Path]
) -> list[tuple[str, Path]]:
def check_holdout_overlap(output_path: Path, resolved_new_files: list[Path]) -> list[tuple[str, Path]]:
"""Return (other_manifest_name, file) pairs where new files clash with existing manifests.
The check is triggered when output_path is (or will be) holdout.manifest, or when a
@@ -641,9 +609,7 @@ def _run_bump(
if gen is None:
new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
else:
new_dirs, log_line = plan_bump_schema(
root_path, kind, gen, reason, by, date, to
)
new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date, to)
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
print("new directories:")
scripts/check_migration_v02_v03.py
+5 -16
View File
@@ -84,19 +84,12 @@ def main() -> int:
mode = model_config.get("mode", "flow")
routed = bool((model_config.get("router") or {}).get("enabled"))
print(f"checkpoint: {args.checkpoint}")
print(
f" mode={mode!r} conditioning={model_config.get('conditioning')!r} "
f"routed={routed} ema={args.ema}"
)
print(f" mode={mode!r} conditioning={model_config.get('conditioning')!r} routed={routed} ema={args.ema}")
stage1_key = "model_ema" if args.ema and "model_ema" in ckpt else "model"
stage2_key = (
"sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
)
stage2_key = "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
if args.ema and stage1_key == "model":
print(
" warning: --ema requested but no model_ema in checkpoint, using raw weights"
)
print(" warning: --ema requested but no model_ema in checkpoint, using raw weights")
# --- old side: the frozen v0.2 snapshot, loaded with the checkpoint's own weights ---
old_stage1, old_stage2 = legacy.build_models(model_config)
@@ -119,9 +112,7 @@ def main() -> int:
print("PASS (construction only, routed checkpoint)")
return 0
remapped1, remapped2 = net.migrate_legacy_state_dict(
ckpt[stage1_key], ckpt[stage2_key]
)
remapped1, remapped2 = net.migrate_legacy_state_dict(ckpt[stage1_key], ckpt[stage2_key])
missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
if missing1 or unexpected1 or missing2 or unexpected2:
@@ -132,9 +123,7 @@ def main() -> int:
new_stage1.eval()
new_stage2.eval()
cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(
model_config, args.batch, args.seed
)
cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(model_config, args.batch, args.seed)
ok = True
with torch.no_grad():
scripts/create_root_files.py
+7 -31
View File
@@ -64,9 +64,7 @@ def parse_detector_spec(spec: str) -> tuple[str, str | None]:
if ":" in spec:
label, config = spec.split(":", 1)
if not label or not config:
raise PlanError(
f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG"
)
raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
return label, config
return spec, None
@@ -94,9 +92,7 @@ def plan_jobs(
raise PlanError(f"--gen must look like 'genN', got {gen!r}")
gen_dir = dataset_root / "raw" / kind / gen
if not gen_dir.is_dir():
raise PlanError(
f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first"
)
raise PlanError(f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first")
jobs = []
for spec in detector_specs:
@@ -124,9 +120,7 @@ def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
return zlib.crc32(key.encode()) & 0x7FFFFFFF
def build_cmd(
executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None
) -> list[str]:
def build_cmd(executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None) -> list[str]:
"""minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
cmd = [str(executable)]
if job.config:
@@ -147,10 +141,7 @@ def run_job(
gen: str,
tmp_root: Path,
) -> JobResult:
workdir = (
tmp_root
/ f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
)
workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
workdir.mkdir(parents=True)
cmd = build_cmd(executable, job, events_per_file, energy_gev)
@@ -174,20 +165,12 @@ def run_job(
job,
False,
None,
f"expected exactly one .root output in {workdir}, found {len(produced)}: "
f"{[p.name for p in produced]}",
f"expected exactly one .root output in {workdir}, found {len(produced)}: {[p.name for p in produced]}",
result.stdout,
result.stderr,
)
dest = (
dataset_root
/ "raw"
/ kind
/ gen
/ job.detector
/ f"shard-{job.shard_index:03d}.root"
)
dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
if dest.exists():
return JobResult(
job,
@@ -279,14 +262,7 @@ def run_make_root(
print(f"executable: {executable}")
for job in planned_jobs:
cmd = build_cmd(executable, job, events_per_file, energy_gev)
dest = (
dataset_root_path
/ "raw"
/ kind
/ gen
/ job.detector
/ f"shard-{job.shard_index:03d}.root"
)
dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
seed = job_seed(kind, gen, job, energy_gev)
print(f" MINICALOSIM_SEED={seed} {' '.join(cmd)} -> {dest}")
scripts/dwarf.py
+40 -116
View File
@@ -80,8 +80,7 @@ def convert(
typer.Option(
"--output",
"-o",
help="Output Parquet file (default: <input>.parquet). Only valid "
"with a single input file and --jobs 1.",
help="Output Parquet file (default: <input>.parquet). Only valid with a single input file and --jobs 1.",
),
] = None,
batch_size: Annotated[
@@ -91,9 +90,7 @@ def convert(
help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
),
] = "100 MB",
tree: Annotated[
str, typer.Option("--tree", help="Tree name inside the ROOT file")
] = "Steps",
tree: Annotated[str, typer.Option("--tree", help="Tree name inside the ROOT file")] = "Steps",
compression: Annotated[
Compression, typer.Option("--compression", help="Parquet compression codec")
] = Compression.snappy,
@@ -129,15 +126,11 @@ def convert(
raise typer.Exit(1)
_warn_if_exceeds_shared_quota(jobs, "--jobs")
compression_value = (
"uncompressed" if compression is Compression.none else compression.value
)
compression_value = "uncompressed" if compression is Compression.none else compression.value
if jobs == 1:
if output is not None and len(root_files) > 1:
typer.echo(
"error: --output can only be used with a single input file", err=True
)
typer.echo("error: --output can only be used with a single input file", err=True)
raise typer.Exit(1)
total_orphaned = 0
for root_file in root_files:
@@ -150,10 +143,7 @@ def convert(
)
total_orphaned += n_orphaned
if total_orphaned:
typer.echo(
f"\n{total_orphaned} orphaned child track(s) dropped across "
f"{len(root_files)} file(s)."
)
typer.echo(f"\n{total_orphaned} orphaned child track(s) dropped across {len(root_files)} file(s).")
return
if output is not None:
@@ -176,9 +166,7 @@ def convert(
@app.command()
def migrate(
root: Annotated[
Path, typer.Argument(help="Dataset root to migrate in place")
] = _DATASET_ROOT_DEFAULT,
root: Annotated[Path, typer.Argument(help="Dataset root to migrate in place")] = _DATASET_ROOT_DEFAULT,
execute: Annotated[
bool,
typer.Option(
@@ -190,8 +178,7 @@ def migrate(
bool,
typer.Option(
"--copy",
help="Copy instead of move, leaving the originals in place "
"(e.g. if another process is still reading them)",
help="Copy instead of move, leaving the originals in place (e.g. if another process is still reading them)",
),
] = False,
) -> None:
@@ -202,12 +189,8 @@ def migrate(
@app.command("bump-gen")
def bump_gen(
reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
by: Annotated[
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
] = None,
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
date: Annotated[
Optional[str],
typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -220,12 +203,8 @@ def bump_gen(
help="Target gen tag (default: one past the current highest)",
),
] = None,
execute: Annotated[
bool, typer.Option("--execute", help="Apply (default: dry run)")
] = False,
root: Annotated[
Path, typer.Option("--root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new raw generation."""
run_bump_gen(
@@ -243,12 +222,8 @@ def bump_gen(
def bump_schema(
gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
by: Annotated[
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
] = None,
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
date: Annotated[
Optional[str],
typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -261,12 +236,8 @@ def bump_schema(
help="Target schema tag (default: one past the current highest)",
),
] = None,
execute: Annotated[
bool, typer.Option("--execute", help="Apply (default: dry run)")
] = False,
root: Annotated[
Path, typer.Option("--root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new schema within a gen."""
run_bump_schema(
@@ -283,9 +254,7 @@ def bump_schema(
@app.command()
def status(
root: Annotated[
Path, typer.Option("--root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""List existing gens/schemas per kind."""
run_status(str(root))
@@ -293,9 +262,7 @@ def status(
@app.command("update-manifest")
def update_manifest(
manifests: Annotated[
list[Path], typer.Argument(help="One or more .manifest files to update")
],
manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
schema: Annotated[
Optional[str],
typer.Option(
@@ -306,9 +273,7 @@ def update_manifest(
] = None,
gen: Annotated[
Optional[str],
typer.Option(
"--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"
),
typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
] = None,
execute: Annotated[
bool,
@@ -316,9 +281,7 @@ def update_manifest(
] = False,
) -> None:
"""Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
run_update_manifest(
[str(m) for m in manifests], schema=schema, execute=execute, gen=gen
)
run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)
@app.command("create-manifest")
@@ -333,22 +296,15 @@ def create_manifest(
typer.Option(
"--pool",
metavar="DETECTOR",
help="Detector name; combined with --type and --root to form "
"<root>/pools/<detector>/<type>.manifest",
help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
),
] = None,
type_: Annotated[
Optional[PoolType],
typer.Option(
"--type", help="Pool type — full, holdout, or dev (required with --pool)"
),
typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
] = None,
root: Annotated[
Path, typer.Option("--root", help="Dataset root (used with --pool)")
] = _DATASET_ROOT_DEFAULT,
execute: Annotated[
bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
] = False,
root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
execute: Annotated[bool, typer.Option("--execute", help="Write the manifest (default: dry run)")] = False,
force: Annotated[
bool,
typer.Option("--force", help="Overwrite the manifest if it already exists"),
@@ -368,9 +324,7 @@ def create_manifest(
@app.command("make-root")
def make_root(
executable: Annotated[
Path, typer.Option("--executable", help="Built minicalosim run_* executable")
],
executable: Annotated[Path, typer.Option("--executable", help="Built minicalosim run_* executable")],
detector: Annotated[
list[str],
typer.Option(
@@ -382,15 +336,9 @@ def make_root(
"Repeatable.",
),
],
num_files: Annotated[
int, typer.Option("--num-files", help="New shards to create per detector")
],
events_per_file: Annotated[
int, typer.Option("--events-per-file", help="nEvents passed to the executable")
],
gen: Annotated[
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
],
num_files: Annotated[int, typer.Option("--num-files", help="New shards to create per detector")],
events_per_file: Annotated[int, typer.Option("--events-per-file", help="nEvents passed to the executable")],
gen: Annotated[str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")],
energy_gev: Annotated[
float | None,
typer.Option(
@@ -401,20 +349,12 @@ def make_root(
"to name the dataset accordingly.",
),
] = None,
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
dataset_root: Annotated[
Path, typer.Option("--dataset-root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
jobs: Annotated[
int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
] = 4,
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
dataset_root: Annotated[Path, typer.Option("--dataset-root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
jobs: Annotated[int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")] = 4,
execute: Annotated[
bool,
typer.Option(
"--execute", help="Actually run jobs (default: dry run / print plan)"
),
typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)"),
] = False,
) -> None:
"""Generate new ROOT shards via a minicalosim executable."""
@@ -441,9 +381,7 @@ class OracleMethod(str, Enum):
@app.command("build-geometry-oracle")
def build_geometry_oracle(
data: Annotated[
Path, typer.Argument(help="Steps parquet file or directory of steps files")
],
data: Annotated[Path, typer.Argument(help="Steps parquet file or directory of steps files")],
out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
method: Annotated[
OracleMethod,
@@ -456,9 +394,7 @@ def build_geometry_oracle(
),
),
] = OracleMethod.slab,
k: Annotated[
int, typer.Option("--k", help="Neighbours for the knn classifier")
] = 1,
k: Annotated[int, typer.Option("--k", help="Neighbours for the knn classifier")] = 1,
subsample: Annotated[
int,
typer.Option("--subsample", help="Max reference points sampled from the data"),
@@ -504,9 +440,7 @@ def build_geometry_oracle(
def warm_cache(
data: Annotated[
Path,
typer.Argument(
help="Parquet file, directory, or .manifest — same as `giant train`'s"
),
typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
],
val_fraction: Annotated[
float,
@@ -518,16 +452,13 @@ def warm_cache(
] = 0.1,
seed: Annotated[
int,
typer.Option(
"--seed", "-s", help="Must match the `giant train` run(s) to warm for"
),
typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
] = 0,
particle_conditioning: Annotated[
Conditioning,
typer.Option(
"--particle-conditioning",
help="Must match the `giant train` run(s)' conditioning.particle.type "
"to warm for",
help="Must match the `giant train` run(s)' conditioning.particle.type to warm for",
),
] = Conditioning.physical,
material_conditioning: Annotated[
@@ -543,21 +474,14 @@ def warm_cache(
bool,
typer.Option(
"--router/--no-router",
help="Warm the process vocabulary too (only takes effect with "
"--router-type process)",
help="Warm the process vocabulary too (only takes effect with --router-type process)",
),
] = False,
router_type: Annotated[
str, typer.Option("--router-type", help="Router implementation name")
] = "energy",
n_experts: Annotated[
int, typer.Option("--n-experts", help="Number of routed experts")
] = 4,
router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
rebuild: Annotated[
bool,
typer.Option(
"--rebuild", help="Ignore any existing sidecar and recompute every section"
),
typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
] = False,
) -> None:
"""Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
scripts/geometry_oracle.py
+1 -3
View File
@@ -38,9 +38,7 @@ def run_build_geometry_oracle(
n_bins=n_bins,
)
print(
f"method: {method} reference points: {oracle.metadata['n_reference_points']:,}"
)
print(f"method: {method} reference points: {oracle.metadata['n_reference_points']:,}")
print("classes (material, layer_id):")
for material, layer_id in oracle.classes:
print(f" {material:<12} layer_id={layer_id}")
scripts/hparam_scan.py
+3 -10
View File
@@ -154,9 +154,7 @@ def run_hparam_scan(
wall_time_s = time.monotonic() - start
if metrics_path.exists():
epochs_completed, final_val_loss, best_val_loss = final_metrics(
metrics_path
)
epochs_completed, final_val_loss, best_val_loss = final_metrics(metrics_path)
append_summary(
summary_path,
{
@@ -171,11 +169,6 @@ def run_hparam_scan(
"wall_time_s": round(wall_time_s, 1),
},
)
print(
f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} "
f"({wall_time_s:.1f}s)"
)
print(f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} ({wall_time_s:.1f}s)")
else:
print(
f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
)
print(f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log")
scripts/migrate_geant_steps.py
+8 -48
View File
@@ -50,12 +50,8 @@ PREDICTED_RE = re.compile(
r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)"
r"_predicted(?P<local>_local)?\.parquet$"
)
SHARD_RE = re.compile(
r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$"
)
LEGACY_PREDICTED_RE = re.compile(
r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$"
)
SHARD_RE = re.compile(r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$")
LEGACY_PREDICTED_RE = re.compile(r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$")
LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?P<ext>root|parquet)$")
@@ -110,24 +106,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
if m:
detector, shard, ext = m["detector"], int(m["shard"]), m["ext"]
if ext == "root":
dst = (
src_root
/ "raw"
/ "steps"
/ GEN
/ detector
/ f"shard-{shard:03d}.root"
)
dst = src_root / "raw" / "steps" / GEN / detector / f"shard-{shard:03d}.root"
else:
dst = (
src_root
/ "processed"
/ "steps"
/ GEN
/ SCHEMA
/ detector
/ f"shard-{shard:03d}.parquet"
)
dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
moves.append((path, dst))
continue
@@ -135,19 +116,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
if m:
ext = m["ext"]
if ext == "root":
dst = (
src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
)
dst = src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
else:
dst = (
src_root
/ "processed"
/ "hits"
/ LEGACY_GEN
/ LEGACY_SCHEMA
/ "pbwo4"
/ "shard-000.parquet"
)
dst = src_root / "processed" / "hits" / LEGACY_GEN / LEGACY_SCHEMA / "pbwo4" / "shard-000.parquet"
moves.append((path, dst))
continue
@@ -165,20 +136,9 @@ def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
for pool, shards in rules.items():
manifest_path = manifest_dir / f"{pool}{MANIFEST_SUFFIX}"
for shard in shards:
dst = (
src_root
/ "processed"
/ "steps"
/ GEN
/ SCHEMA
/ detector
/ f"shard-{shard:03d}.parquet"
)
dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
manifests[manifest_path].append((shard, dst))
return {
k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)]
for k, v in manifests.items()
}
return {k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)] for k, v in manifests.items()}
def run_migration(root: str, execute: bool, copy: bool) -> None:
scripts/profile_analysis_costs.py
+2 -6
View File
@@ -94,9 +94,7 @@ def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
"post_dx": post_dir[:, 0],
"post_dy": post_dir[:, 1],
"post_dz": post_dir[:, 2],
"edep": np.where(
is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep
),
"edep": np.where(is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep),
"step_length": np.where(is_synthetic, 0.0, step_length),
"material": rng.choice(_MATERIALS, size=n),
"layer_id": rng.integers(0, 30, size=n),
@@ -163,9 +161,7 @@ def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
)
def _time(
spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path
) -> float:
def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path) -> float:
t0 = time.perf_counter()
compute_reduced(
spec_id,
scripts/steps_to_parquet_parallel.py
+4 -17
View File
@@ -49,9 +49,7 @@ def latest_schema_tag(processed_gen_dir: Path) -> str | None:
return best_tag
def resolve_destination(
root_file: Path, dataset_root: Path, schema_override: str | None
) -> Path:
def resolve_destination(root_file: Path, dataset_root: Path, schema_override: str | None) -> Path:
"""Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.
@@ -66,12 +64,7 @@ def resolve_destination(
raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")
parts = rel.parts
if (
len(parts) != 5
or parts[0] != "raw"
or not GEN_RE.match(parts[2])
or not parts[4].endswith(".root")
):
if len(parts) != 5 or parts[0] != "raw" or not GEN_RE.match(parts[2]) or not parts[4].endswith(".root"):
raise DestinationError(
f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
f"under {dataset_root} (got relative path: {rel})"
@@ -216,13 +209,7 @@ def run_parallel_job(
print(f" {root_file}", file=sys.stderr)
raise SystemExit(1)
total_orphaned = sum(
int(m.group(1))
for _, _, stdout, _ in results
for m in _ORPHAN_RE.finditer(stdout)
)
total_orphaned = sum(int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout))
if total_orphaned:
print(
f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)."
)
print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
print(f"\nAll {len(results)} conversion(s) completed.")
tests/legacy/network_v02_snapshot.py
+36 -131
View File
@@ -37,11 +37,7 @@ class SinusoidalEmbedding(nn.Module):
super().__init__()
assert dim % 2 == 0, "dim must be even"
half = dim // 2
freqs = torch.exp(
-math.log(10000)
* torch.arange(half, dtype=torch.float32)
/ max(half - 1, 1)
)
freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
self.register_buffer("freqs", freqs)
def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -107,9 +103,7 @@ class ConditionEncoder(nn.Module):
pdg_e = self.pdg_emb(cond_cat[:, 0])
mat_e = self.mat_emb(cond_cat[:, 1])
else:
particle_phys = cond_cont[
:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
]
particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
pdg_e = self.particle_mlp(particle_phys)
mat_e = self.material_mlp(material_phys)
@@ -168,12 +162,7 @@ class DenoisingMLP(nn.Module):
)
merged_cond_dim = time_dim + cond_out_dim
self.input_proj = nn.Linear(x_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_proj = nn.Linear(hidden_dim, x_dim)
# Predicts n_sec as classification over {0, 1, ..., k_max}.
# Applied to the condition encoding (not the diffused latent).
@@ -286,12 +275,7 @@ class SecondaryDecoder(nn.Module):
)
merged_cond_dim = time_dim + cond_out_dim
self.input_proj = nn.Linear(sec_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_proj = nn.Linear(hidden_dim, sec_dim)
def forward(
@@ -345,12 +329,7 @@ class WGANGenerator(nn.Module):
conditioning=conditioning,
)
self.input_proj = nn.Linear(noise_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_proj = nn.Linear(hidden_dim, x_dim)
self.n_sec_head = nn.Sequential(
nn.Linear(cond_out_dim, hidden_dim // 2),
@@ -410,12 +389,7 @@ class Critic(nn.Module):
conditioning=conditioning,
)
self.input_proj = nn.Linear(x_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_norm = nn.LayerNorm(hidden_dim)
self.out_proj = nn.Linear(hidden_dim, 1)
@@ -466,12 +440,7 @@ class WGANSecondaryGenerator(nn.Module):
conditioning=conditioning,
)
self.input_proj = nn.Linear(noise_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_proj = nn.Linear(hidden_dim, sec_dim)
def forward(
@@ -515,12 +484,7 @@ class SecondaryCritic(nn.Module):
conditioning=conditioning,
)
self.input_proj = nn.Linear(sec_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_norm = nn.LayerNorm(hidden_dim)
self.out_proj = nn.Linear(hidden_dim, 1)
@@ -550,9 +514,7 @@ class Router(nn.Module):
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
raise NotImplementedError
def combine_weights(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
probs = self.gate(cond_cont, cond_cat)
if not (self.gumbel and self.training):
return probs
@@ -562,26 +524,18 @@ class Router(nn.Module):
def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
return self.gate(cond_cont, cond_cat).argmax(dim=-1)
def balance_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,)
return (importance.std() / (importance.mean() + 1e-8)) ** 2
def classify_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
return torch.zeros((), device=cond_cont.device)
def entropy_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
return norm_entropy
def gate_stats(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
gate = self.gate(cond_cont, cond_cat) # (B, n_experts)
row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1) # (B,)
norm_entropy = row_entropy.mean() / math.log(self.n_experts)
@@ -602,9 +556,7 @@ def register_router(name: str):
def build_router(name: str, n_experts: int, **kwargs) -> Router:
if name not in ROUTER_REGISTRY:
raise ValueError(
f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
)
raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
cls = ROUTER_REGISTRY[name]
accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
filtered = {k: v for k, v in kwargs.items() if k in accepted}
@@ -644,8 +596,7 @@ class EnergyRouter(Router):
if learn_width or learn_temperature:
if not (width_min_ratio < 1.0 < width_max_ratio):
raise ValueError(
f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
f"({width_max_ratio}) must bracket 1.0"
f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
)
self._width_lo = width_min_ratio * temperature
self._width_hi = width_max_ratio * temperature
@@ -658,10 +609,7 @@ class EnergyRouter(Router):
centers = torch.linspace(-2.0, 2.0, n_experts)
else:
if len(centers_init) != n_experts:
raise ValueError(
f"centers_init has {len(centers_init)} values, "
f"expected n_experts={n_experts}"
)
raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
centers = torch.tensor(list(centers_init), dtype=torch.float32)
if learn_centers:
self.centers = nn.Parameter(centers)
@@ -702,9 +650,7 @@ class PdgRouter(Router):
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim)
d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
-1
) # (B, n_experts)
d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1) # (B, n_experts)
return torch.softmax(-d2 / self.temperature, dim=-1)
@@ -736,9 +682,7 @@ class ProcessRouter(Router):
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
def classify_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
@@ -756,14 +700,10 @@ class ComposedRouter(Router):
joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0)
for router in self.routers[1:]:
g = router.gate(cond_cont, cond_cat) # (B, n_i)
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
1
) # (B, prod so far)
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1) # (B, prod so far)
return joint
def classify_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
total = torch.zeros((), device=cond_cont.device)
for router in self.routers:
total = total + router.classify_loss(cond_cont, cond_cat, labels)
@@ -796,12 +736,7 @@ class ExpertTrunk(nn.Module):
) -> None:
super().__init__()
self.input_proj = nn.Linear(in_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
self.out_proj = nn.Linear(hidden_dim, in_dim)
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
@@ -891,9 +826,7 @@ class RoutedDenoisingMLP(nn.Module):
t_emb = self.time_emb(t)
c_emb = self.cond_enc(cond_cont, cond_cat)
cond = torch.cat([t_emb, c_emb], dim=-1)
return _route_forward(
self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
)
return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)
def predict_n_sec(
self,
@@ -957,9 +890,7 @@ class RoutedSecondaryDecoder(nn.Module):
t_emb = self.time_emb(t)
c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
cond = torch.cat([t_emb, c_emb], dim=-1)
return _route_forward(
self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
)
return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)
_STAGE1_MODEL_KEYS = {
@@ -1006,9 +937,7 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]:
_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
def _check_router_conditioning_compat(
router_types: list[str], conditioning: str
) -> None:
def _check_router_conditioning_compat(router_types: list[str], conditioning: str) -> None:
bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
if bad and conditioning == "physical":
raise ValueError(
@@ -1019,9 +948,7 @@ def _check_router_conditioning_compat(
)
def _build_router_from_cfg(
router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding"
) -> Router:
def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding") -> Router:
shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
if router_cfg["type"] == "composed":
axes = _parse_composed_axes(router_cfg)
@@ -1030,9 +957,7 @@ def _build_router_from_cfg(
router.gumbel = bool(router_cfg.get("gumbel", False))
return router
_check_router_conditioning_compat([router_cfg["type"]], conditioning)
router_kwargs = {
k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
}
router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
router_kwargs.setdefault("pdg_vocab", pdg_vocab)
router_kwargs.setdefault("mat_vocab", mat_vocab)
router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
@@ -1042,15 +967,9 @@ def _build_router_from_cfg(
def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
if model_config.get("mode") == "wgan":
stage1 = WGANGenerator(
**{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
)
stage1 = WGANGenerator(**{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS})
sec_decoder = WGANSecondaryGenerator(
**{
k: v
for k, v in model_config.items()
if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
}
**{k: v for k, v in model_config.items() if k in _WGAN_SEC_GENERATOR_MODEL_KEYS}
)
return stage1, sec_decoder
@@ -1061,44 +980,30 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
shared = dict(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
expert_hidden_dim=model_config.get("expert_hidden_dim")
or model_config.get("hidden_dim", 128),
expert_n_blocks=model_config.get("expert_n_blocks")
or model_config.get("n_blocks", 3),
expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
expert_n_blocks=model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
emb_dim=model_config.get("emb_dim", EMB_DIM),
dropout=model_config.get("dropout", 0.1),
conditioning=model_config.get("conditioning", "embedding"),
)
conditioning = shared["conditioning"]
stage1 = RoutedDenoisingMLP(
router=_build_router_from_cfg(
router_cfg, pdg_vocab, mat_vocab, conditioning
),
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
k_max=model_config.get("k_max", K_MAX),
**shared,
)
sec_decoder = RoutedSecondaryDecoder(
router=_build_router_from_cfg(
router_cfg, pdg_vocab, mat_vocab, conditioning
),
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
**shared,
)
return stage1, sec_decoder
stage1 = DenoisingMLP(
**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
)
sec_decoder = SecondaryDecoder(
**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
)
stage1 = DenoisingMLP(**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS})
sec_decoder = SecondaryDecoder(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
return stage1, sec_decoder
def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
critic = Critic(
**{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
)
sec_critic = SecondaryCritic(
**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
)
critic = Critic(**{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS})
sec_critic = SecondaryCritic(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
return critic, sec_critic
tests/test_bump_dataset_version.py
+17 -105
View File
@@ -22,9 +22,7 @@ def test_git_user_name_returns_none_on_timeout(monkeypatch):
def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
dirs, log_line = plan_bump_gen(
tmp_path, "steps", "first generation", None, "2026-01-01"
)
dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
assert dirs == [
tmp_path / "raw" / "steps" / "gen1",
tmp_path / "processed" / "steps" / "gen1" / "schema1",
@@ -56,9 +54,7 @@ def test_bump_gen_kinds_are_independent(tmp_path):
def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
dirs, log_line = plan_bump_schema(
tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01"
)
dirs, log_line = plan_bump_schema(tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01")
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"]
assert "`gen1`/`schema1`" in log_line
@@ -66,18 +62,14 @@ def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
def test_bump_schema_increments_within_its_gen(tmp_path):
(tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
(tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
dirs, _ = plan_bump_schema(
tmp_path, "steps", "gen1", "next schema", None, "2026-01-01"
)
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]
def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
(tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
(tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
dirs, _ = plan_bump_schema(
tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01"
)
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
@@ -91,9 +83,7 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
def test_bump_gen_to_specific_tag(tmp_path):
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
dirs, log_line = plan_bump_gen(
tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5"
)
dirs, log_line = plan_bump_gen(tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5")
assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
assert "`gen5`" in log_line
@@ -125,9 +115,7 @@ def test_bump_schema_to_specific_tag(tmp_path):
def test_bump_schema_rejects_invalid_to_tag(tmp_path):
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
try:
plan_bump_schema(
tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
)
plan_bump_schema(tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3")
assert False, "expected SystemExit"
except SystemExit:
pass
@@ -164,15 +152,7 @@ def _make_parquet(path):
def test_update_manifest_bumps_to_specified_schema(tmp_path):
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema2"
/ "pbwo4"
/ "shard-000.parquet"
)
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -194,15 +174,7 @@ def test_update_manifest_auto_detects_highest_schema(tmp_path):
for schema in ("schema1", "schema2", "schema3"):
d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
d.mkdir(parents=True)
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema3"
/ "pbwo4"
/ "shard-000.parquet"
)
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
parquet.touch()
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -231,15 +203,7 @@ def test_update_manifest_reports_missing_targets(tmp_path):
def test_update_manifest_skips_already_at_target(tmp_path):
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema2"
/ "pbwo4"
/ "shard-000.parquet"
)
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -254,15 +218,7 @@ def test_update_manifest_skips_already_at_target(tmp_path):
def test_update_manifest_preserves_comments_and_blanks(tmp_path):
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema2"
/ "pbwo4"
/ "shard-000.parquet"
)
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -278,15 +234,7 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
def test_update_manifest_bumps_gen(tmp_path):
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen2"
/ "schema1"
/ "pbwo4"
/ "shard-000.parquet"
)
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "pbwo4" / "shard-000.parquet"
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -303,15 +251,7 @@ def test_update_manifest_bumps_gen(tmp_path):
def test_update_manifest_bumps_gen_and_schema(tmp_path):
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen2"
/ "schema3"
/ "pbwo4"
/ "shard-000.parquet"
)
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "pbwo4" / "shard-000.parquet"
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -328,15 +268,7 @@ def test_update_manifest_bumps_gen_and_schema(tmp_path):
def test_apply_update_manifest_writes_file(tmp_path):
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema2"
/ "pbwo4"
/ "shard-000.parquet"
)
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -358,24 +290,8 @@ def test_apply_update_manifest_writes_file(tmp_path):
def test_create_manifest_writes_relative_paths(tmp_path):
pq1 = (
tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema2"
/ "pbwo4"
/ "shard-000.parquet"
)
pq2 = (
tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema2"
/ "pbwo4"
/ "shard-001.parquet"
)
pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
_make_parquet(pq1)
_make_parquet(pq2)
@@ -419,9 +335,7 @@ def test_run_create_manifest_refuses_to_overwrite_existing_output(tmp_path):
output.write_text("original contents\n")
try:
bump_dataset_version.run_create_manifest(
[str(pq)], execute=True, output=str(output)
)
bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output))
assert False, "expected SystemExit"
except SystemExit:
pass
@@ -435,9 +349,7 @@ def test_run_create_manifest_force_overwrites_existing_output(tmp_path):
output.parent.mkdir(parents=True)
output.write_text("original contents\n")
bump_dataset_version.run_create_manifest(
[str(pq)], execute=True, output=str(output), force=True
)
bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output), force=True)
assert output.read_text() != "original contents\n"
tests/test_catalog.py
+2 -7
View File
@@ -13,9 +13,7 @@ from tests.test_analysis_reduce import _reference_frame, _rollout_frame
def _build_ctx() -> Context:
r, t = _rollout_frame(), _reference_frame()
return build_context(
r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
)
return build_context(r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
@pytest.fixture(scope="module")
@@ -142,10 +140,7 @@ def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
# 4 chunks over only 2 distinct event_ids also exercises empty chunks.
n_chunks = 4 if spec.chunkable else 1
parts = [
spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks)))
for k in range(n_chunks)
]
parts = [spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
chunked = spec.finalize(parts, ctx)
assert chunked.id == unchunked.id
tests/test_cli_new_run.py
+1 -3
View File
@@ -101,9 +101,7 @@ def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
assert "already has last.pt" in result.output
assert not (out_dir / "config.toml").exists()
result = runner.invoke(
app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"]
)
result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"])
assert result.exit_code == 0, result.output
assert (out_dir / "config.toml").exists()
tests/test_cli_predict.py
+3 -9
View File
@@ -116,9 +116,7 @@ def test_ref_yaml_includes_comment_when_provided(tmp_path):
dataset = tmp_path / "full.manifest"
pred_uuid = str(uuid.uuid4())
ref_path = _write_prediction_ref(
checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3"
)
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3")
data = yaml.safe_load(ref_path.read_text())
assert data["comment"] == "baseline sweep run 3"
@@ -133,9 +131,7 @@ def test_ref_timestamp_is_iso_format(tmp_path):
checkpoint.touch()
pred_uuid = str(uuid.uuid4())
ref_path = _write_prediction_ref(
checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
)
ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
data = yaml.safe_load(ref_path.read_text())
# Must parse without error and be timezone-aware (UTC).
@@ -150,9 +146,7 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
checkpoint.touch()
pred_uuid = str(uuid.uuid4())
ref_path = _write_prediction_ref(
checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
)
ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
data = yaml.safe_load(ref_path.read_text())
assert data["checkpoint"].startswith("/")
tests/test_condor.py
+8 -27
View File
@@ -62,9 +62,7 @@ def _fake_venv(repo_dir: Path) -> None:
giant.chmod(0o755)
def _prep(
rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1
) -> Path:
def _prep(rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1) -> Path:
"""``prep`` with small test-sized context bins/sampling."""
return prep(
rollout_yaml,
@@ -224,16 +222,12 @@ def test_write_submit_description(tmp_path: Path):
assert "--chunk" in body and "--run-dir" in body
def test_write_submit_requires_synced_venv(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
run_dir = _prep(_write_inputs(tmp_path))
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
# No `giant` next to the (fake) active interpreter, so this falls through
# to repo_dir/.venv/bin/giant, which _write_inputs/_prep also didn't create.
monkeypatch.setattr(
sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python")
)
monkeypatch.setattr(sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python"))
with pytest.raises(FileNotFoundError, match="uv sync"):
write_submit(cfg)
@@ -241,9 +235,7 @@ def test_write_submit_requires_synced_venv(
def test_write_submit_remote_flag(tmp_path: Path):
run_dir = _prep(_write_inputs(tmp_path))
_fake_venv(tmp_path)
cfg = SubmitConfig(
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True
)
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True)
txt = write_submit(cfg).read_text()
assert "+RemoteJob = True" in txt
assert "ProvidesETPResources" not in txt
@@ -253,9 +245,7 @@ def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
assert get_spec("router_gating").chunkable is False
run_dir = _prep(_write_inputs(tmp_path), chunks=4)
_fake_venv(tmp_path)
cfg = SubmitConfig(
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
)
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
write_submit(cfg)
jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
counts: dict[str, int] = {}
@@ -272,9 +262,7 @@ def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
_job_walltimes instead of a clear error here."""
run_dir = _prep(_write_inputs(tmp_path), chunks=2)
_fake_venv(tmp_path)
cfg = SubmitConfig(
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
)
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
with pytest.raises(ValueError, match="n_chunks"):
write_submit(cfg)
@@ -297,16 +285,9 @@ def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
run_dir = _prep(_write_inputs(tmp_path), chunks=2)
meta = RunMeta.load(run_dir / "run_meta.json")
_fake_venv(tmp_path)
cfg = SubmitConfig(
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2
)
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2)
write_submit(cfg)
jobs = {
(i, int(k)): int(w)
for i, k, w in (
line.split(",") for line in (run_dir / "jobs.txt").read_text().split()
)
}
jobs = {(i, int(k)): int(w) for i, k, w in (line.split(",") for line in (run_dir / "jobs.txt").read_text().split())}
for chunk in range(2):
expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
assert jobs[("marginal_edep", chunk)] == expected
tests/test_config.py
+18 -57
View File
@@ -112,9 +112,7 @@ def test_migrate_config_lambda_nsec_and_lambda_s2():
def test_migrate_config_wgan_knobs_map_to_both_stages():
new = gconfig.migrate_config(
{"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}}
)
new = gconfig.migrate_config({"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}})
for stage in ("stage1_model", "stage2_model"):
assert new[stage]["wgan"]["n_critic"] == 3
assert new[stage]["wgan"]["gp_weight"] == 5.0
@@ -122,9 +120,7 @@ def test_migrate_config_wgan_knobs_map_to_both_stages():
def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
new = gconfig.migrate_config(
{"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}}
)
new = gconfig.migrate_config({"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}})
for stage in ("stage1_model", "stage2_model"):
assert new[stage]["hidden_dim"] == 128
assert new[stage]["n_res_blocks"] == 4
@@ -132,9 +128,7 @@ def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes():
new = gconfig.migrate_config(
{"model": {"emb_dim": 32, "conditioning": "embedding"}}
)
new = gconfig.migrate_config({"model": {"emb_dim": 32, "conditioning": "embedding"}})
for axis in ("particle", "material"):
assert new["conditioning"][axis]["emb_dim"] == 32
assert new["conditioning"][axis]["type"] == "embedding"
@@ -181,11 +175,7 @@ def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false():
def test_migrate_config_router_nonzero_expert_dims_raises():
cfg = {
"model": {
"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}
}
}
cfg = {"model": {"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}}}
try:
gconfig.migrate_config(cfg)
assert False, "expected ValueError"
@@ -267,9 +257,7 @@ def test_merge_cli_overrides_nested_override_keeps_siblings():
assert cfg["stage1_model"]["hidden_dim"] == 256 # untouched sibling section
def test_merge_cli_overrides_file_then_explicit_override_precedence(
tmp_path, monkeypatch
):
def test_merge_cli_overrides_file_then_explicit_override_precedence(tmp_path, monkeypatch):
monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
path = tmp_path / "config.toml"
_write_toml(
@@ -317,9 +305,7 @@ def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, c
assert "current999" in captured.err
def test_merge_cli_overrides_no_warning_on_matching_git_hash(
tmp_path, monkeypatch, capsys
):
def test_merge_cli_overrides_no_warning_on_matching_git_hash(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
path = tmp_path / "config.toml"
_write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n")
@@ -328,9 +314,7 @@ def test_merge_cli_overrides_no_warning_on_matching_git_hash(
assert capsys.readouterr().err == ""
def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
tmp_path, monkeypatch, capsys
):
def test_merge_cli_overrides_no_warning_when_git_hash_unknown(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
path = tmp_path / "config.toml"
_write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n")
@@ -339,9 +323,7 @@ def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
assert capsys.readouterr().err == ""
def test_merge_cli_overrides_no_warning_when_meta_section_absent(
tmp_path, monkeypatch, capsys
):
def test_merge_cli_overrides_no_warning_when_meta_section_absent(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
path = tmp_path / "config.toml"
path.write_text("[train]\nepochs = 5\n")
@@ -351,12 +333,8 @@ def test_merge_cli_overrides_no_warning_when_meta_section_absent(
def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch):
monkeypatch.setattr(
gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243"
)
cfg = gconfig.merge_cli_overrides(
gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {}
)
monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {})
assert cfg["stage1_model"]["generator"] == "flow"
assert cfg["stage1_model"]["hidden_dim"] == 256
assert cfg["stage2_model"]["hidden_dim"] == 256
@@ -431,10 +409,7 @@ def _cfg_with(**dotted_overrides):
def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
assert (
gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW)
== "20260729_1430"
)
assert gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) == "20260729_1430"
def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix():
@@ -499,9 +474,7 @@ def test_default_out_dir_name_router_gumbel_shown_when_enabled():
"stage1_model.router.gumbel": True,
}
)
assert (
gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
)
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
@@ -524,9 +497,7 @@ def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
# First 6 by priority: stage1_generator, stage2_generator, stage2_decoder,
# stage2_history, particle_type_target, stage1_router — stage2_router
# overflows into the hash suffix.
assert name.startswith(
"20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+"
)
assert name.startswith("20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+")
def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
@@ -758,15 +729,11 @@ def test_validate_config_ar_checks_skipped_under_one_shot():
# ---------------------------------------------------------------------------
def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
tmp_path, monkeypatch, capsys
):
def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
ckpt_path = tmp_path / "best.pt"
ckpt_path.write_bytes(b"")
_write_toml(
tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n"
)
_write_toml(tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n")
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
@@ -776,9 +743,7 @@ def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
assert "current999" in captured.err
def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
tmp_path, monkeypatch, capsys
):
def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
ckpt_path = tmp_path / "best.pt"
ckpt_path.write_bytes(b"")
@@ -787,15 +752,11 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
assert capsys.readouterr().err == ""
def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
tmp_path, monkeypatch, capsys
):
def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
ckpt_path = tmp_path / "best.pt"
ckpt_path.write_bytes(b"")
_write_toml(
tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n"
)
_write_toml(tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n")
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
assert capsys.readouterr().err == ""
tests/test_create_root_files.py
+5 -15
View File
@@ -16,9 +16,7 @@ SimJob = create_root_files.SimJob
PlanError = create_root_files.PlanError
def _write_fake_executable(
path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0
) -> Path:
def _write_fake_executable(path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0) -> Path:
"""Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files
into its own cwd (so callers can verify each job gets an isolated workdir
and that the workdir ends up holding *only* the .root output, matching
@@ -89,17 +87,13 @@ def test_next_shard_index_continues_past_existing(tmp_path):
def test_plan_jobs_rejects_missing_gen(tmp_path):
with pytest.raises(PlanError):
plan_jobs(
["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
)
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")
def test_plan_jobs_rejects_malformed_gen(tmp_path):
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
with pytest.raises(PlanError):
plan_jobs(
["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
)
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")
def test_plan_jobs_continues_from_existing_shards(tmp_path):
@@ -108,9 +102,7 @@ def test_plan_jobs_continues_from_existing_shards(tmp_path):
(gen_dir / "pbwo4" / "shard-000.root").touch()
(gen_dir / "pbwo4" / "shard-001.root").touch()
jobs = plan_jobs(
["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
)
jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")
assert [j.shard_index for j in jobs] == [2, 3, 4]
assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
@@ -320,9 +312,7 @@ def test_run_all_caps_concurrency(tmp_path):
assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}
intervals = [json.loads(d.read_text()) for d in dests]
events = sorted(
[(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
)
events = sorted([(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals])
concurrent = 0
peak = 0
for _, delta in events:
tests/test_dwarf.py
+1 -3
View File
@@ -51,9 +51,7 @@ def test_convert_rejects_output_with_multiple_files(tmp_path):
def test_convert_rejects_output_with_parallel_jobs(tmp_path):
root_file = tmp_path / "shard.root"
root_file.touch()
result = runner.invoke(
app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
)
result = runner.invoke(app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"])
assert result.exit_code != 0
assert "--output cannot be combined with --jobs > 1" in result.output
tests/test_geometry.py
+5 -15
View File
@@ -44,9 +44,7 @@ def test_iter_point_batches_without_post_columns_yields_pre_only(tmp_path):
(pos, mat, lay) = next(g._iter_point_batches(path))
assert pos.shape == (5, 3)
np.testing.assert_allclose(
pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32)
)
np.testing.assert_allclose(pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
assert list(mat) == ["G4_AIR"] * 5
np.testing.assert_array_equal(lay, np.arange(5))
@@ -63,12 +61,8 @@ def test_iter_point_batches_with_post_columns_doubles_and_concatenates_points(
# Every step contributes both its pre_pos and post_pos, sharing the
# step's material/layer_id label — so batches double in length.
assert pos.shape == (10, 3)
np.testing.assert_allclose(
pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32)
)
np.testing.assert_allclose(
pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32)
)
np.testing.assert_allclose(pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
np.testing.assert_allclose(pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32))
assert list(mat) == ["G4_AIR"] * 10
np.testing.assert_array_equal(lay, np.concatenate([np.arange(5), np.arange(5)]))
@@ -208,9 +202,7 @@ def test_slab_classes_discovered():
def test_slab_query_labels_by_depth():
orc = _build_slab()
pos = np.array(
[[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
) # layer 0, gap, layer 1
pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]) # layer 0, gap, layer 1
material, layer_id, escaped = orc.query(pos)
assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
assert list(layer_id) == [0, -1, 1]
@@ -239,9 +231,7 @@ def test_slab_save_load_roundtrip(tmp_path):
orc.save(p)
loaded = g.GeometryOracle.load(p)
pos = np.array(
[[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
)
pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]])
m0, l0, e0 = orc.query(pos)
m1, l1, e1 = loaded.query(pos)
assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
tests/test_loader.py
+2 -6
View File
@@ -315,9 +315,7 @@ def test_build_index_maps_from_files_ordering_independent_of_file_order(tmp_path
def test_build_index_maps_from_files_numeric_sort_for_nuclear_codes(tmp_path):
path = tmp_path / "a.parquet"
pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(
path
)
pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(path)
pdg_map, _ = build_index_maps_from_files([path])
assert list(pdg_map.keys()) == [11, 22, 1000060120]
@@ -398,9 +396,7 @@ def test_load_event_ids_applies_offset(tmp_path):
path = tmp_path / "a.parquet"
pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path)
offset = event_id_offset(1)
np.testing.assert_array_equal(
load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2]
)
np.testing.assert_array_equal(load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2])
def test_load_event_ids_raises_when_event_id_reaches_stride(tmp_path):
tests/test_materials.py
+1 -5
View File
@@ -23,11 +23,7 @@ def test_get_material_properties_unfilled_entry_raises():
def test_get_material_properties_returns_filled_entry_from_injected_table():
table = {
"G4_Pb": MaterialProperties(
z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59
)
}
table = {"G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59)}
props = get_material_properties("G4_Pb", table)
assert props.z_eff == 82.0
assert props.a_eff == 207.2
tests/test_migration_v02_v03.py
+3 -9
View File
@@ -56,9 +56,7 @@ def _random_batch(seed: int):
def _assert_bit_identical(a: torch.Tensor, b: torch.Tensor, label: str) -> None:
assert a.shape == b.shape, f"{label}: shape mismatch {a.shape} vs {b.shape}"
assert torch.equal(a, b), (
f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
)
assert torch.equal(a, b), f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
def _run_migration_check(mode: str, conditioning: str) -> None:
@@ -137,9 +135,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
assert new_stage1.n_sec_head is not None
assert new_stage2.n_sec_head is None
remapped1, remapped2 = net.migrate_legacy_state_dict(
old_stage1.state_dict(), old_stage2.state_dict()
)
remapped1, remapped2 = net.migrate_legacy_state_dict(old_stage1.state_dict(), old_stage2.state_dict())
missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
assert not missing1 and not unexpected1
@@ -159,9 +155,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t)
_assert_bit_identical(old_out1, new_out1, f"stage1 output ({mode}, {conditioning})")
_assert_bit_identical(
old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})"
)
_assert_bit_identical(old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})")
_assert_bit_identical(old_out2, new_out2, f"stage2 output ({mode}, {conditioning})")
tests/test_network.py
+13 -42
View File
@@ -85,9 +85,7 @@ def test_stage1_model_gradients_flow():
def test_stage1_model_no_n_sec_head_by_default():
"""Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
it moves to stage 2."""
model = Stage1Model(
pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
)
model = Stage1Model(pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG)
assert model.n_sec_head is None
@@ -121,29 +119,18 @@ def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
k_max = 15
assert (
stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16)
== k_max * SEC_SLOT_DIM
)
assert (
stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16)
== k_max * SEC_SLOT_DIM
)
assert stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
assert stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
k_max = 15
assert stage2_trunk_sec_dim(
{"target": "onehot"}, "wgan", k_max, emb_dim=16
) == k_max * (CONT_SLOT_DIM + 16)
assert stage2_trunk_sec_dim({"target": "onehot"}, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)
def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
k_max = 15
assert (
stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16)
== k_max * CONT_SLOT_DIM
)
assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
# --- ConditionEncoder onehot mode -------------------------------------------
@@ -440,16 +427,12 @@ def test_stage2_autoregressive_history_invalid_raises():
@pytest.mark.parametrize("history", ["markov", "attention"])
def test_stage2_autoregressive_forward_shape(target, generator, history):
B, K, emb_dim = 4, 5, 6
model = _build_stage2_ar(
target, generator, emb_dim=emb_dim, k_max=K, history=history
)
model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K, history=history)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
stage1_out = torch.randn(B, 9)
type_dim = stage2_type_dim({"target": target}, emb_dim)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
B, K, CONT_SLOT_DIM + type_dim
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
if generator == "wgan":
x_t = torch.randn(B, K, model.noise_dim)
@@ -488,9 +471,7 @@ def test_stage2_autoregressive_predict_type_shape():
cond_cat = torch.zeros(B, 2, dtype=torch.long)
stage1_out = torch.randn(B, 9)
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
B, K, CONT_SLOT_DIM + type_dim
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
out = model.predict_type(
cond_cont,
cond_cat,
@@ -511,9 +492,7 @@ def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, gen
cond_cat = torch.zeros(B, 2, dtype=torch.long)
stage1_out = torch.randn(B, 9)
type_dim = stage2_type_dim({"target": target}, emb_dim)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
B, K, CONT_SLOT_DIM + type_dim
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
with pytest.raises(RuntimeError):
model.predict_type(
cond_cont,
@@ -533,9 +512,7 @@ def test_stage2_autoregressive_gradients_flow_wgan_onehot():
cond_cat = torch.zeros(B, 2, dtype=torch.long)
stage1_out = torch.randn(B, 9)
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
B, K, CONT_SLOT_DIM + type_dim
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
z = torch.randn(B, K, model.noise_dim)
gen_out = model(
z,
@@ -560,9 +537,7 @@ def test_stage2_autoregressive_gradients_flow_onehot():
cond_cat = torch.zeros(B, 2, dtype=torch.long)
stage1_out = torch.randn(B, 9)
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
B, K, CONT_SLOT_DIM + type_dim
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
x_t = torch.randn(B, K, token_dim)
t = torch.rand(B, K)
@@ -601,17 +576,13 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
itself rather than `AttentionHistory` in isolation
(`test_attention_history_step_matches_forward` covers that lower layer)."""
B, K, emb_dim = 3, 6, 6
model = _build_stage2_ar(
"physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention"
)
model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention")
model.eval()
type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
hist_in_dim = CONT_SLOT_DIM + type_dim
own_feat = torch.randn(B, K, hist_in_dim) # token i's own raw feature
has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
history_feat = torch.cat(
[torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1
)
history_feat = torch.cat([torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1)
with torch.no_grad():
expected = model.history_encoder(history_feat, has_prev_full)
tests/test_particles.py
+2 -6
View File
@@ -49,9 +49,7 @@ def test_ground_state_nucleus_resolved_via_particle_package():
"""He-4 (Z=2, A=4) is a common nuclide in `particle`'s ground-state table."""
mass, charge = particle_mass_charge(1000020040)
assert charge == pytest.approx(2.0)
assert mass == pytest.approx(
4 * 931.494, rel=0.05
) # near A*amu, binding-energy-corrected
assert mass == pytest.approx(4 * 931.494, rel=0.05) # near A*amu, binding-energy-corrected
def test_nuclear_isomer_falls_back_to_z_a_decode():
@@ -174,9 +172,7 @@ def test_decode_topn_class_other_drop_returns_zero_sentinel():
def test_decode_topn_class_other_sample_stays_within_members():
topn_map, n_classes = _topn_fixture()
rng = np.random.default_rng(0)
out = decode_topn_class(
np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng
)
out = decode_topn_class(np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng)
assert set(out.tolist()) <= {2212, 2112}
tests/test_phase2.py
+20 -60
View File
@@ -57,9 +57,7 @@ def _sec_decoder(pdg=3, mat=2, conditioning="embedding"):
def _cond(B=8, pdg=3, mat=2):
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
)
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
return cond_cont, cond_cat
@@ -154,9 +152,7 @@ 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
@@ -169,9 +165,7 @@ 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)
@@ -182,9 +176,7 @@ def test_flow_matching_loss_secondary_has_grad():
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)
flow_matching_loss_secondary(
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
).backward()
flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask).backward()
assert any(p.grad is not None for p in decoder.parameters())
@@ -220,9 +212,7 @@ def test_flow_matching_loss_secondary_ar_scalar():
x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
sec_mask = torch.ones(B, K, dtype=torch.bool)
loss = flow_matching_loss_secondary_ar(
decoder,
@@ -246,9 +236,7 @@ def test_flow_matching_loss_secondary_ar_mask_zeros_padding():
x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
sec_mask = torch.zeros(B, K, dtype=torch.bool)
loss = flow_matching_loss_secondary_ar(
decoder,
@@ -271,9 +259,7 @@ def test_flow_matching_loss_secondary_ar_has_grad():
x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
)
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
sec_mask = torch.ones(B, K, dtype=torch.bool)
flow_matching_loss_secondary_ar(
decoder,
@@ -299,9 +285,7 @@ def test_sample_secondaries_shapes():
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
sec_cont, sec_phys, sec_valid = sample_secondaries(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3
)
sec_cont, sec_phys, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3)
assert sec_cont.shape == (B, K_MAX, 4)
assert sec_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM)
assert sec_valid.shape == (B, K_MAX)
@@ -314,9 +298,7 @@ def test_sample_secondaries_valid_mask_matches_n_sec():
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
_, _, sec_valid = sample_secondaries(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
)
_, _, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
for i, n in enumerate(n_sec_pred.tolist()):
assert sec_valid[i, :n].all()
assert not sec_valid[i, n:].any()
@@ -350,9 +332,7 @@ def test_encode_secondaries_energy_conservation():
pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
sec_cont = encode_secondaries(
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
)
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
assert sec_cont.shape == (N, K_MAX, 6)
assert np.isfinite(sec_cont).all()
@@ -399,9 +379,7 @@ def test_encode_secondaries_stick_logits_match_naive_reference():
logit = np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP)
expected[row, i] = logit
np.testing.assert_allclose(
stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4
)
np.testing.assert_allclose(stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4)
def test_encode_secondaries_direction_encoding():
@@ -513,9 +491,7 @@ def test_encode_secondaries_physical_columns_match_ground_truth_pdg():
sec_valid[0, 0] = True
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
sec_cont = encode_secondaries(
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
)
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
mass, charge = particle_mass_charge(11)
assert sec_cont[0, 0, 4] == pytest.approx(log_transform(np.array([mass]))[0])
assert sec_cont[0, 0, 5] == pytest.approx(charge)
@@ -549,9 +525,7 @@ def test_decode_secondaries_valid_slots_sum_to_e_sec():
e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
sec_cont, n_sec, e_sec, pre_dir
)
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
valid_sum = (sec_E * sec_valid).sum(axis=1)
has_secondaries = n_sec > 0
@@ -574,9 +548,7 @@ def test_decode_secondaries_zero_n_sec_has_zero_energy():
e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
sec_cont, n_sec, e_sec, pre_dir
)
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
assert not sec_valid.any()
np.testing.assert_allclose(sec_E, 0.0)
@@ -596,9 +568,7 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
sec_cont, n_sec, e_sec, pre_dir
)
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
for i, k in enumerate(n_sec):
if k == 0:
@@ -622,12 +592,8 @@ def test_decode_secondaries_rescale_preserves_relative_shares():
n_sec = np.array([4])
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
sec_E_small, _, _, _, sec_valid = decode_secondaries(
sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir
)
sec_E_large, _, _, _, _ = decode_secondaries(
sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir
)
sec_E_small, _, _, _, sec_valid = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
sec_E_large, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir)
ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
@@ -649,20 +615,14 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
sec_valid[0, 0] = True
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
sec_cont = encode_secondaries(
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
)
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
norm = Normalizer()
norm.mean = np.array([-2.0, 0.5], dtype=np.float32)
norm.std = np.array([3.0, 1.5], dtype=np.float32)
sec_cont_normed = sec_cont.copy()
sec_cont_normed[:, :, 4:6] = norm.transform(
sec_cont[:, :, 4:6].reshape(-1, 2)
).reshape(N, K_MAX, 2)
sec_cont_normed[:, :, 4:6] = norm.transform(sec_cont[:, :, 4:6].reshape(-1, 2)).reshape(N, K_MAX, 2)
n_sec = np.array([1])
_, _, sec_mass, sec_charge, _ = decode_secondaries(
sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm
)
_, _, sec_mass, sec_charge, _ = decode_secondaries(sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm)
assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
tests/test_pipeline.py
+9 -27
View File
@@ -48,9 +48,7 @@ def _make_synthetic_steps(path, n_events=20, seed=0):
pre_dir = np.array([0.0, 0.0, 1.0])
post_dir = _unit(rng.normal(size=3))
post_pos = pre_pos + step_length * pre_dir
sec_energies = (
list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
)
sec_energies = list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)]
sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)]
rows.append(
@@ -211,9 +209,7 @@ def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
@pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
tmp_path, data, monkeypatch
):
def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(tmp_path, data, monkeypatch):
# num_workers>0 makes DataLoader actually fork worker subprocesses
# (unlike every other test here, which runs with num_workers=0) — pytest
# itself is multi-threaded, hence Python's fork-safety warning below.
@@ -223,9 +219,7 @@ def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
@pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
def test_run_train_job_no_warning_when_num_workers_within_shared_quota(
tmp_path, data, monkeypatch
):
def test_run_train_job_no_warning_when_num_workers_within_shared_quota(tmp_path, data, monkeypatch):
monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2
echo = _run(data, tmp_path / "out", num_workers=2)
assert not any("exceeds" in m for m in echo)
@@ -299,12 +293,8 @@ def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path,
from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
assert cond_norm.mean is not None and cond_norm.std is not None
np.testing.assert_allclose(
cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0
)
np.testing.assert_allclose(
cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0
)
np.testing.assert_allclose(cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0)
np.testing.assert_allclose(cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0)
material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
@@ -342,12 +332,8 @@ def test_run_train_job_matches_uncached_output(tmp_path, data):
cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False)
for key in ("cond", "target", "sec_phys"):
np.testing.assert_allclose(
uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]
)
np.testing.assert_allclose(
uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]
)
np.testing.assert_allclose(uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"])
np.testing.assert_allclose(uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"])
assert uncached["pdg_map"] == cached["pdg_map"]
assert uncached["mat_map"] == cached["mat_map"]
@@ -376,9 +362,7 @@ def test_seed_energy_router_falls_back_to_default_and_warns_when_no_samples():
router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
cond_norm = _fitted_cond_norm()
echoed = []
_seed_energy_router(
router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append
)
_seed_energy_router(router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append)
assert "centers_init" not in router_cfg
assert len(echoed) == 1
assert "falls back to default centers" in echoed[0]
@@ -393,9 +377,7 @@ def test_seed_energy_router_seeds_centers_from_data_quantiles():
# units as the conditioning column being normalized against.
energy_quantiles = np.linspace(1.0, 10.0, 33).astype(np.float32)
echoed = []
_seed_energy_router(
router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append
)
_seed_energy_router(router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append)
assert "centers_init" in router_cfg
centers = np.asarray(router_cfg["centers_init"], dtype=np.float32)
tests/test_render.py
+7 -16
View File
@@ -73,10 +73,7 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
"x",
{
"edges": [0, 1, 2],
"groups": {
lbl: {"rollout": [1, 2], "reference": [2, 1]}
for lbl in ("a", "b", "c", "d")
},
"groups": {lbl: {"rollout": [1, 2], "reference": [2, 1]} for lbl in ("a", "b", "c", "d")},
"log_y": True,
},
),
@@ -126,9 +123,7 @@ def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
for r in reduced:
r.save(tmp_path / "reduced" / f"{r.id}.json")
try:
render_mod.render_all(
tmp_path / "reduced", tmp_path / "plots", run_gallery=True
)
render_mod.render_all(tmp_path / "reduced", tmp_path / "plots", run_gallery=True)
except RuntimeError as e:
pytest.skip(f"LaTeX rendering unavailable: {e}")
@@ -145,9 +140,7 @@ def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkey
(run_dir / "reduced").mkdir(parents=True)
merge_calls = []
monkeypatch.setattr(
condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd))
)
monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
meta = condor_mod.RunMeta(
rollout="rollout.parquet",
reference="reference.parquet",
@@ -157,9 +150,9 @@ def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkey
)
monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))
Reduced(
"s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]}
).save(run_dir / "reduced" / "s.json")
Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]}).save(
run_dir / "reduced" / "s.json"
)
try:
pdfs = render_mod.render_run(run_dir)
@@ -363,9 +356,7 @@ def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps():
def test_plot_metadata_includes_note_and_run_meta_parameters():
r = Reduced(
"u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"}
)
r = Reduced("u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"})
meta = render_mod._plot_metadata(r, {"title": "run-1", "checkpoint": "ckpt.pt"})
assert meta["note"] == "no router data"
assert meta["parameters"] == {"checkpoint": "ckpt.pt"}
tests/test_rollout.py
+6 -18
View File
@@ -121,12 +121,8 @@ def fake_material_props(monkeypatch):
import giant.materials as gm
fake = {
"G4_AIR": gm.MaterialProperties(
z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5
),
"G4_PbWO4": gm.MaterialProperties(
z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
),
"G4_AIR": gm.MaterialProperties(z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5),
"G4_PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7),
}
monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
return fake
@@ -496,9 +492,7 @@ def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
@pytest.mark.parametrize("generator2", ["flow", "wgan"])
def test_rollout_physical_target_decoder_generator_matrix(
fake_material_props, decoder, generator2
):
def test_rollout_physical_target_decoder_generator_matrix(fake_material_props, decoder, generator2):
"""Every (decoder, stage2 generator) combination under
particle_type.target="physical" must run to completion and conserve
energy."""
@@ -542,9 +536,7 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
assert len(rec["event_id"]) > 0
# Every spawned secondary's nominal pdg must be one decode_topn_class can
# actually produce (the topn map's known classes + its "other" members).
possible = set(PDG_TOPN_MAP.class_map.keys()) | set(
PDG_TOPN_MAP.other_members.keys()
)
possible = set(PDG_TOPN_MAP.class_map.keys()) | set(PDG_TOPN_MAP.other_members.keys())
secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
assert secondary_pdgs <= possible
@@ -589,9 +581,7 @@ def _onehot_conditioning_models():
return s1.eval(), s2.eval()
def _run_onehot_conditioning(
pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP
):
def _run_onehot_conditioning(pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP):
s1, s2 = _onehot_conditioning_models()
cond, tgt, sec_phys = _norms()
return rollout(
@@ -643,9 +633,7 @@ def test_rollout_embedding_target_end_to_end(decoder):
conditioning's own particle embedding table, so every resolved PDG must
be a real member of the dense training vocab (pdg_map) — unlike
"onehot", there is no "other" bucket to fall outside of."""
s1, s2 = _models_v3(
conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4
)
s1, s2 = _models_v3(conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4)
rec = _run_v3(s1, s2, conditioning="embedding")
assert len(rec["event_id"]) > 0
secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
tests/test_router.py
+18 -54
View File
@@ -25,9 +25,7 @@ MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
def _cond(B=8, pdg=3, mat=2):
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
)
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
return cond_cont, cond_cat
@@ -78,9 +76,7 @@ def test_energy_router_gate_partition_of_unity():
def test_energy_router_top1_matches_gate_argmax():
router = EnergyRouter(n_experts=4)
cond_cont, cond_cat = _cond(16)
assert torch.equal(
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
)
assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
def test_energy_router_hardens_as_temperature_shrinks():
@@ -128,12 +124,8 @@ def test_energy_router_centers_init_wrong_length_raises():
def test_energy_router_centers_init_respects_learn_centers_flag():
learned = EnergyRouter(
n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True
)
fixed = EnergyRouter(
n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False
)
learned = EnergyRouter(n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True)
fixed = EnergyRouter(n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False)
assert isinstance(learned.centers, torch.nn.Parameter)
assert not isinstance(fixed.centers, torch.nn.Parameter)
@@ -160,9 +152,7 @@ def test_energy_router_learn_width_matches_fixed_temperature_at_init():
per-expert width must reproduce the fixed-temperature gate exactly."""
centers_init = [-1.0, 0.0, 0.5, 1.5]
fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
learned = EnergyRouter(
n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True
)
learned = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True)
cond_cont, cond_cat = _cond(16)
torch.testing.assert_close(
learned.gate(cond_cont, cond_cat),
@@ -195,21 +185,15 @@ def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True)
except ValueError:
return
raise AssertionError(
"expected ValueError for learn_width and learn_temperature both set"
)
raise AssertionError("expected ValueError for learn_width and learn_temperature both set")
def test_energy_router_width_ratio_bounds_must_bracket_one_raises():
try:
EnergyRouter(
n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0
)
EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0)
except ValueError:
return
raise AssertionError(
"expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0"
)
raise AssertionError("expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0")
def test_energy_router_effective_width_stays_within_bounds():
@@ -246,9 +230,7 @@ def test_energy_router_learn_width_hardens_when_pushed_to_floor():
"""Pushing every expert's width toward the (tiny) floor should harden the
gate to a one-hot at the nearest center, generalizing the fixed-
temperature->0 hardening test to the per-expert path."""
router = EnergyRouter(
n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0
)
router = EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0)
with torch.no_grad():
router.raw_width.fill_(-1e6)
cond_cont, cond_cat = _cond(16)
@@ -264,9 +246,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
expert's own gate share, without needing to touch any other expert's
width — the "each expert learns its own coverage independently" property
this feature is meant to add."""
router = EnergyRouter(
n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]
)
router = EnergyRouter(n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0])
cond_cont, cond_cat = _cond(4)
cond_cont[:, 3] = 3.0 # fixed energy, unequal distance to each center
@@ -280,9 +260,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
def test_build_router_threads_learn_width_kwargs_through():
router = build_router(
"energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0
)
router = build_router("energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0)
assert isinstance(router, EnergyRouter)
assert router.learn_width is True
assert isinstance(router.raw_width, torch.nn.Parameter)
@@ -419,9 +397,7 @@ def test_pdg_router_gate_partition_of_unity():
def test_pdg_router_top1_matches_gate_argmax():
router = PdgRouter(n_experts=4, pdg_vocab=3)
cond_cont, cond_cat = _cond(16)
assert torch.equal(
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
)
assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
def test_pdg_router_hardens_as_temperature_shrinks():
@@ -586,9 +562,7 @@ def test_process_router_gate_partition_of_unity():
def test_process_router_top1_matches_gate_argmax():
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
cond_cont, cond_cat = _cond(16)
assert torch.equal(
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
)
assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
def test_process_router_balance_loss_is_nonnegative_scalar():
@@ -663,16 +637,12 @@ def test_build_models_routed_with_process_router():
def test_composed_router_n_experts_is_product():
router = ComposedRouter(
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
)
router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
assert router.n_experts == 12
def test_composed_router_gate_partition_of_unity():
router = ComposedRouter(
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
)
router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
cond_cont, cond_cat = _cond(16, pdg=5)
g = router.gate(cond_cont, cond_cat)
assert g.shape == (16, 12)
@@ -709,9 +679,7 @@ def test_composed_router_top1_factors_into_per_axis_argmax():
def test_composed_router_supports_different_expert_counts_per_axis():
router = ComposedRouter(
[EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
)
router = ComposedRouter([EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)])
assert router.n_experts == 10
cond_cont, cond_cat = _cond(8, pdg=5)
assert router.gate(cond_cont, cond_cat).shape == (8, 10)
@@ -719,9 +687,7 @@ def test_composed_router_supports_different_expert_counts_per_axis():
def test_composed_router_classify_loss_sums_sub_router_losses():
"""energy/pdg both default to zero, so the composed loss should too."""
router = ComposedRouter(
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
)
router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
cond_cont, cond_cat = _cond(16, pdg=5)
labels = torch.randint(0, 4, (16,))
loss = router.classify_loss(cond_cont, cond_cat, labels)
@@ -964,9 +930,7 @@ def test_routed_stage1_eval_dispatch_matches_manual_grouping():
idx = model.trunk.router.top1(cond_cont, cond_cat)
manual = torch.zeros_like(x_t)
for i in range(B):
manual[i] = model.trunk.experts[int(idx[i])](
x_t[i : i + 1], cond[i : i + 1]
)[0]
manual[i] = model.trunk.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
tests/test_sample.py
+10 -30
View File
@@ -30,9 +30,7 @@ def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]
def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
)
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
return cond_cont, cond_cat
@@ -43,12 +41,8 @@ def _conditioning_for(target: str) -> str:
return "embedding" if target == "embedding" else "physical"
def _stage2_oneshot(
target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2
) -> Stage2OneShot:
particle_cfg, material_cfg = _particle_material_cfg(
_conditioning_for(target), emb_dim
)
def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot:
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
particle_type_cfg = {"target": target}
# build_models (giant/model/network.py) computes sec_dim this same way
# before constructing Stage2OneShot — its own default (SEC_DIM, the
@@ -78,9 +72,7 @@ def _stage2_ar(
k_max: int = 5,
history: str = "markov",
) -> Stage2Autoregressive:
particle_cfg, material_cfg = _particle_material_cfg(
_conditioning_for(target), emb_dim
)
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
return Stage2Autoregressive(
pdg_vocab=pdg,
mat_vocab=mat,
@@ -163,9 +155,7 @@ def test_sample_secondaries_flow_shapes_by_target(target):
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
sec_cont, sec_type, sec_valid = sample_secondaries(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
)
sec_cont, sec_type, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
assert sec_valid.shape == (B, K_MAX)
@@ -180,9 +170,7 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
sec_cont, sec_type, sec_valid = sample_secondaries_wgan(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
)
sec_cont, sec_type, sec_valid = sample_secondaries_wgan(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
assert sec_valid.shape == (B, K_MAX)
@@ -196,15 +184,11 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_ar_shapes(target, generator, history):
B, k_max, emb_dim = 4, 5, 6
decoder = _stage2_ar(
target, generator, emb_dim=emb_dim, k_max=k_max, history=history
)
decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max, history=history)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, k_max + 1, (B,))
sec_cont, sec_type, sec_valid = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
)
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
assert sec_cont.shape == (B, k_max, CONT_SLOT_DIM)
assert sec_type.shape == (B, k_max, _expected_type_dim(target, emb_dim))
assert sec_valid.shape == (B, k_max)
@@ -220,9 +204,7 @@ def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 2, k_max])
_, _, sec_valid = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
)
_, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
for i, n in enumerate(n_sec_pred.tolist()):
assert sec_valid[i, :n].all()
assert not sec_valid[i, n:].any()
@@ -237,8 +219,6 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 1, 1])
sec_cont, sec_type, sec_valid = sample_secondaries_ar(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
)
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
assert sec_valid.tolist() == [[False], [True], [True]]
tests/test_setup_cache.py
+4 -13
View File
@@ -14,9 +14,7 @@ from giant.data.transforms import Normalizer
def _touch_parquet(path, n=1):
path.parent.mkdir(parents=True, exist_ok=True)
pd.DataFrame(
{"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}
).to_parquet(path)
pd.DataFrame({"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}).to_parquet(path)
return path
@@ -29,9 +27,7 @@ def _normalizer(width=3):
def _entry(n_train_steps=100, sample=None):
sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) if sample is None else sample
return NormalizerEntry(
_normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample
)
return NormalizerEntry(_normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample)
# ── sidecar_path ─────────────────────────────────────────────────────────
@@ -39,9 +35,7 @@ def _entry(n_train_steps=100, sample=None):
def test_sidecar_path_single_file(tmp_path):
f = tmp_path / "shard.parquet"
assert (
setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
)
assert setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
def test_sidecar_path_directory(tmp_path):
@@ -51,10 +45,7 @@ def test_sidecar_path_directory(tmp_path):
def test_sidecar_path_manifest(tmp_path):
m = tmp_path / "pools" / "full.manifest"
assert (
setup_cache.sidecar_path(m)
== tmp_path / "pools" / "full.manifest.giant_train_cache.json"
)
assert setup_cache.sidecar_path(m) == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
# ── fingerprint_files ────────────────────────────────────────────────────
tests/test_steps_to_parquet.py
+3 -9
View File
@@ -24,18 +24,14 @@ def _frame() -> pl.DataFrame:
def test_e_sec_sums_child_first_step_energy():
out, n_orphaned = steps_to_parquet._add_secondary_attributes(_frame())
assert n_orphaned == 0
e_sec = dict(
zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
)
e_sec = dict(zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"]))
assert e_sec[(1, 0, 0)] == 15.0 # one child, first-step pre_E 15
assert e_sec[(1, 0, 1)] == 50.0 # two children, 20 + 30
def test_e_sec_zero_when_no_children():
out, _ = steps_to_parquet._add_secondary_attributes(_frame())
childless = out.filter(
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
)
childless = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1))
assert childless["e_sec"].item() == 0.0
@@ -69,9 +65,7 @@ def test_orphaned_child_track_is_dropped_not_nulled():
}
)
out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
row = out.filter(
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)
)
row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
assert n_orphaned == 1
assert row["child_track_ids"].to_list() == [[2]]
tests/test_steps_to_parquet_parallel.py
+2 -20
View File
@@ -124,31 +124,13 @@ def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
def test_resolve_destination_uses_latest_schema(tmp_path):
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
dest = resolve_destination(root_file, tmp_path, schema_override=None)
assert (
dest
== tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema3"
/ "pbwo4"
/ "shard-000.parquet"
)
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
def test_resolve_destination_schema_override_wins(tmp_path):
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
assert (
dest
== tmp_path
/ "processed"
/ "steps"
/ "gen1"
/ "schema9"
/ "pbwo4"
/ "shard-000.parquet"
)
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"
def test_resolve_destination_errors_without_any_schema(tmp_path):
tests/test_train.py
+21 -63
View File
@@ -74,9 +74,7 @@ def test_wandb_run_config_includes_full_cfg_and_param_counts():
"stage1_model": {"generator": "flow"},
"stage2_model": {"generator": "wgan"},
}
wcfg = _wandb_run_config(
cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100}
)
wcfg = _wandb_run_config(cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100})
assert wcfg["train"] == {"lr": 3e-4}
assert wcfg["stage1_model"] == {"generator": "flow"}
assert wcfg["stage2_model"] == {"generator": "wgan"}
@@ -162,9 +160,7 @@ def test_type_repr_shapes_and_values(target):
expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim
assert repr_.shape == (B, K, expected_width)
if target == "physical":
assert torch.equal(
repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
)
assert torch.equal(repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM])
if target == "onehot":
assert torch.all(repr_.sum(-1) == 1.0)
@@ -180,9 +176,7 @@ def test_type_repr_shapes_and_values(target):
("embedding", "wgan"),
],
)
def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
target, generator
):
def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target, generator):
"""Regression test tying the refactor together: _assemble_stage2_real is
now defined as _assemble_stage2_ar_target(...).flatten(1)."""
B, emb_dim = 4, 6
@@ -192,12 +186,8 @@ def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
if target == "embedding":
cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
particle_type_cfg = {"target": target}
flat = _assemble_stage2_real(
sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
)
unflat = _assemble_stage2_ar_target(
sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
)
flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
assert torch.equal(unflat.flatten(1), flat)
@@ -206,9 +196,7 @@ def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width():
sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
cond_enc = torch.nn.Module()
out = _assemble_stage2_ar_inputs(
sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim
)
out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim)
assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
assert out["has_prev"].shape == (B, K_MAX)
assert out["remaining_frac"].shape == (B, K_MAX)
@@ -221,9 +209,7 @@ def test_relax_onehot_type_slice_grad_probe_populates_both_norms():
B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
grad_probe: dict[str, float] = {}
out = _relax_onehot_type_slice(
x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe
)
out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe)
out.sum().backward()
assert grad_probe["cont"] >= 0.0
assert grad_probe["type"] >= 0.0
@@ -324,9 +310,7 @@ def _fake_batches(n_batches, batch_size, seed=0):
sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
proc_idx = torch.zeros(batch_size, dtype=torch.long)
sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
batches.append(
(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
)
batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
return batches
@@ -409,17 +393,13 @@ def _run_train(cfg, out_dir, resume_path=None):
),
(
"stage2_onehot_target_wgan",
lambda cfg: cfg["stage2_model"].__setitem__(
"particle_type", {"target": "onehot", "lambda": 1.0}
),
lambda cfg: cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
),
(
"stage2_onehot_target_flow",
lambda cfg: (
cfg["stage2_model"].__setitem__("generator", "flow"),
cfg["stage2_model"].__setitem__(
"particle_type", {"target": "onehot", "lambda": 1.0}
),
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
),
),
(
@@ -427,9 +407,7 @@ def _run_train(cfg, out_dir, resume_path=None):
lambda cfg: (
cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
cfg["stage2_model"].__setitem__(
"particle_type", {"target": "embedding", "lambda": 1.0}
),
cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
),
),
(
@@ -438,18 +416,14 @@ def _run_train(cfg, out_dir, resume_path=None):
cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
cfg["stage2_model"].__setitem__("generator", "flow"),
cfg["stage2_model"].__setitem__(
"particle_type", {"target": "embedding", "lambda": 1.0}
),
cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
),
),
(
"ar_wgan_onehot",
lambda cfg: (
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
cfg["stage2_model"].__setitem__(
"particle_type", {"target": "onehot", "lambda": 1.0}
),
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
),
),
(
@@ -461,9 +435,7 @@ def _run_train(cfg, out_dir, resume_path=None):
lambda cfg: (
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
cfg["stage2_model"].__setitem__("generator", "flow"),
cfg["stage2_model"].__setitem__(
"particle_type", {"target": "onehot", "lambda": 1.0}
),
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
),
),
(
@@ -473,9 +445,7 @@ def _run_train(cfg, out_dir, resume_path=None):
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
cfg["stage2_model"].__setitem__("generator", "flow"),
cfg["stage2_model"].__setitem__(
"particle_type", {"target": "embedding", "lambda": 1.0}
),
cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
),
),
(
@@ -491,9 +461,7 @@ def _run_train(cfg, out_dir, resume_path=None):
cfg["stage1_model"].__setitem__("generator", "wgan"),
cfg["stage2_model"].__setitem__("generator", "flow"),
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
cfg["stage2_model"].__setitem__(
"particle_type", {"target": "onehot", "lambda": 1.0}
),
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
),
),
],
@@ -585,9 +553,7 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
ema_decay=0.0,
steps_per_epoch=4,
)
trainer = WGANStageTrainer(
spec, models["stage1"], critics["stage1"], torch.device("cpu")
)
trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu"))
assert trainer.model.n_sec_head is None
batch = _fake_batches(1, 8)[0]
stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -595,9 +561,7 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
spec = StageSpec(
name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0
)
spec = StageSpec(name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0)
with pytest.raises(NotImplementedError):
FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
@@ -608,9 +572,7 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
@pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"])
@pytest.mark.parametrize("history", ["markov", "attention"])
@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
teacher_forcing, history, stage2_generator
):
def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(teacher_forcing, history, stage2_generator):
"""v0.3.0 step 7: history='attention' and teacher_forcing in
{'scheduled', 'never'} must actually train — a stage-2 AR trainer.step()
must run and produce a finite loss, for every {history} x
@@ -629,9 +591,7 @@ def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
model_config = _model_config(cfg)
models = build_models(model_config)
critics = build_critics(model_config)
trainers = build_stage_trainers(
cfg, models, critics, torch.device("cpu"), total_train_batches=4
)
trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
trainer = trainers["stage2"]
batch = _fake_batches(1, 4)[0]
stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -665,9 +625,7 @@ def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
with open(out_dir / "metrics.csv", newline="") as f:
rows = list(csv.DictReader(f))
assert len(rows) == cfg["train"]["epochs"]
loss_col = (
"stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
)
loss_col = "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
assert all(math.isfinite(float(r[loss_col])) for r in rows)
tests/test_transforms.py
+12 -38
View File
@@ -162,9 +162,7 @@ def test_local_frame_rotation_normalizes_non_unit_pre_dir():
post_dir = rng.standard_normal((N, 3)).astype(np.float32)
post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(
np.float32
)
pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(np.float32)
expected = local_frame_rotation(pre_dir_unit, post_dir)
result = local_frame_rotation(pre_dir_scaled, post_dir)
np.testing.assert_allclose(result, expected, atol=1e-4)
@@ -200,14 +198,10 @@ def test_reconstruct_post_pos_straight_line():
step_length = rng.uniform(0.1, 5.0, size=N).astype(np.float32)
post_pos = pre_pos + step_length[:, None] * pre_dir
travel_dir_local = local_frame_rotation(
pre_dir, travel_direction(pre_pos, post_pos)
)
travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
np.testing.assert_allclose(travel_dir_local, np.tile([0, 0, 1], (N, 1)), atol=1e-4)
reconstructed = reconstruct_post_pos(
pre_pos, pre_dir, step_length, travel_dir_local
)
reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
@@ -221,12 +215,8 @@ def test_reconstruct_post_pos_general_roundtrip():
post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
step_length = np.linalg.norm(post_pos - pre_pos, axis=1).astype(np.float32)
travel_dir_local = local_frame_rotation(
pre_dir, travel_direction(pre_pos, post_pos)
)
reconstructed = reconstruct_post_pos(
pre_pos, pre_dir, step_length, travel_dir_local
)
travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
@@ -356,9 +346,7 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
def test_build_features_proc_idx_zero_without_proc_map():
data = _minimal_step_data(
3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
)
data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
*_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map)
@@ -367,9 +355,7 @@ def test_build_features_proc_idx_zero_without_proc_map():
def test_build_features_proc_idx_looks_up_proc_map():
data = _minimal_step_data(
3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
)
data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
proc_map = {"compt": 0, "phot": 1, "eIoni": 2}
@@ -396,9 +382,7 @@ def test_build_features_require_secondaries_ok_when_no_secondaries():
data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
_, _, _, _, sec_cont, *_ = build_features(
data, pdg_map, mat_map, require_secondaries=True
)
_, _, _, _, sec_cont, *_ = build_features(data, pdg_map, mat_map, require_secondaries=True)
assert not sec_cont.any()
@@ -413,11 +397,7 @@ def fake_material_props(monkeypatch):
in by the user (see giant.materials.MaterialPropertiesNotFilledError)."""
import giant.materials as gm
fake = {
"PbWO4": gm.MaterialProperties(
z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
)
}
fake = {"PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7)}
monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
return fake
@@ -455,9 +435,7 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props):
assert cond_cont.shape[1] == COND_DIM
mass, charge = particle_mass_charge(11)
expected_log_mass = log_transform(np.array([mass]))[0]
np.testing.assert_allclose(
cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5
)
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5)
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], charge)
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 2], 75.6) # z_eff
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 3], 205.3) # a_eff
@@ -500,9 +478,7 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
material_conditioning="physical",
)
np.testing.assert_allclose(
cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
)
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])))
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0])
@@ -714,9 +690,7 @@ def test_welford_accumulator_matches_naive_running_mean_reference():
naive_M2 = np.zeros(F)
naive_n = 0
for chunk in chunks:
naive_mean, naive_M2, naive_n = naive_update(
naive_mean, naive_M2, naive_n, chunk
)
naive_mean, naive_M2, naive_n = naive_update(naive_mean, naive_M2, naive_n, chunk)
acc = _WelfordAccumulator(F)
for chunk in chunks:
tests/test_validate.py
+2 -6
View File
@@ -57,9 +57,7 @@ def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int
sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
proc_idx = torch.zeros(B, dtype=torch.long)
sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
batches.append(
(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
)
batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
return batches
@@ -71,9 +69,7 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch
s1, s2 = _tiny_models()
loader = _loader(n_sec_value=0)
def _fake_resolve_n_sec(
stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
):
def _fake_resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred):
return torch.zeros(cond_cont.size(0), dtype=torch.long)
monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)
tests/test_wgan.py
+3 -9
View File
@@ -152,9 +152,7 @@ def test_sample_secondaries_wgan_shape():
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, K_MAX, (B,))
sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(
model, cond_cont, cond_cat, stage1_out, n_sec_pred
)
sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(model, cond_cont, cond_cat, stage1_out, n_sec_pred)
assert sec_cont.shape == (B, K_MAX, 4)
assert sec_phys.shape == (B, K_MAX, 2)
assert sec_valid.shape == (B, K_MAX)
@@ -183,9 +181,7 @@ def test_gradient_penalty_masked():
mask = _mask(B, n_sec)
real = torch.randn(B, SEC_DIM) * mask
fake = torch.randn(B, SEC_DIM) * mask
gp = gradient_penalty(
lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask
)
gp = gradient_penalty(lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask)
assert gp.item() >= 0.0
@@ -195,9 +191,7 @@ def test_critic_loss_scalar_and_grad():
cond_cont, cond_cat = _cond(B)
real = torch.randn(B, X_DIM)
fake = torch.randn(B, X_DIM)
loss = critic_loss(
lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0
)
loss = critic_loss(lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0)
assert loss.shape == ()
loss.backward()
assert any(p.grad is not None for p in critic.parameters())