From 60c2ca1985eb66853dd881d2f6bc441726fd02ed Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 24 Jul 2026 13:32:25 +0200 Subject: [PATCH] analyze: add MoE router gating/share diagnostic plots New "model" family in the gallery: router_gating (mean soft gate weight vs. pre-step energy, showing the router's soft decision boundaries) and router_share_by_pdg/router_share_by_process (stacked top-1 dispatch share by species / true physics process). Needs a live checkpoint's Router, so it's a documented exception to the rest of the package's polars/numpy-only contract; gracefully degrades to a placeholder for non-MoE checkpoints. Co-Authored-By: Claude Sonnet 5 --- giant/analysis/catalog.py | 26 ++- giant/analysis/condor.py | 4 +- giant/analysis/reduced.py | 13 +- giant/analysis/render.py | 77 ++++++++ giant/analysis/router_gating.py | 339 ++++++++++++++++++++++++++++++++ tests/test_catalog.py | 14 ++ tests/test_router_gating.py | 126 ++++++++++++ 7 files changed, 592 insertions(+), 7 deletions(-) create mode 100644 giant/analysis/router_gating.py create mode 100644 tests/test_router_gating.py diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index a91398a..a0e5bc0 100644 --- a/giant/analysis/catalog.py +++ b/giant/analysis/catalog.py @@ -37,6 +37,11 @@ from giant.analysis.reduce import ( weighted_profile, ) from giant.analysis.reduced import Reduced +from giant.analysis.router_gating import ( + compute_router_gating, + compute_router_share_by_pdg, + compute_router_share_by_process, +) from giant.analysis.sources import Side, open_side, physical_steps, secondaries from giant.analysis.variables import RANGED_VARS, cos_scatter_expr @@ -50,9 +55,10 @@ class Bundle: t_all: pl.LazyFrame # reference, all rows r_phys: pl.LazyFrame # rollout, physical steps only t_phys: pl.LazyFrame # reference, physical steps only + checkpoint: str | None = None # from the rollout YAML; router_gating only @classmethod - def open(cls, rollout, reference, ctx: Context) -> "Bundle": + def open(cls, rollout, reference, ctx: Context, checkpoint=None) -> "Bundle": r_all = open_side(rollout, Side.rollout) t_all = open_side(reference, Side.reference) return cls( @@ -61,6 +67,7 @@ class Bundle: t_all=t_all, r_phys=physical_steps(r_all, Side.rollout), t_phys=physical_steps(t_all, Side.reference), + checkpoint=checkpoint, ) @@ -559,6 +566,23 @@ def build_catalog() -> list[PlotSpec]: PlotSpec("sec_count_per_species", "secondaries", _sec_count_per_species), PlotSpec("sec_energy", "secondaries", _sec_energy), PlotSpec("sec_cos_angle", "secondaries", _sec_cos_angle), + PlotSpec( + "router_gating", + "model", + lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys), + ), + PlotSpec( + "router_share_by_pdg", + "model", + lambda b: compute_router_share_by_pdg( + b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs + ), + ), + PlotSpec( + "router_share_by_process", + "model", + lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys), + ), ] return specs diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index e2ef763..b9ad621 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -152,10 +152,11 @@ def compute_reduced( reference: str | Path, shared: str | Path, out: str | Path, + checkpoint: str | None = None, ) -> Path: """Core: run one plot's reduction against explicit paths → ``Reduced`` JSON.""" ctx = Context.load(shared) - bundle = Bundle.open(rollout, reference, ctx) + bundle = Bundle.open(rollout, reference, ctx, checkpoint=checkpoint) reduced = get_spec(spec_id).compute(bundle) out = Path(out) reduced.save(out) @@ -172,6 +173,7 @@ def compute_one(spec_id: str, run_dir: str | Path) -> Path: meta.reference, run_path / "shared.json", run_path / "reduced" / f"{spec_id}.json", + checkpoint=meta.plot_meta.get("checkpoint"), ) diff --git a/giant/analysis/reduced.py b/giant/analysis/reduced.py index 947641e..836c923 100644 --- a/giant/analysis/reduced.py +++ b/giant/analysis/reduced.py @@ -12,11 +12,14 @@ from dataclasses import asdict, dataclass, field from pathlib import Path # Reduced.kind values: -# "overlay_hist" rollout vs reference density histogram over shared edges -# "grouped_hist" one panel per group (energy/pdg/material), each an overlay -# "profile" edep-weighted mean +/- event-RMS vs depth/radius, two series -# "bar" per-category rollout vs reference bars (share / counts) -# "single_hist" one series only (e.g. rollout leakage; reference has none) +# "overlay_hist" rollout vs reference density histogram over shared edges +# "grouped_hist" one panel per group (energy/pdg/material), each an overlay +# "profile" edep-weighted mean +/- event-RMS vs depth/radius, two series +# "bar" per-category rollout vs reference bars (share / counts) +# "single_hist" one series only (e.g. rollout leakage; reference has none) +# "router_gating" stacked mean MoE gate weight vs energy, rollout + reference +# "router_share" stacked bar of MoE top-1 dispatch share by category +# "unavailable" plot not applicable to this run (e.g. non-MoE checkpoint) @dataclass diff --git a/giant/analysis/render.py b/giant/analysis/render.py index d5066c8..35cd3d6 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -137,12 +137,89 @@ def _render_bar(r: Reduced, params: dict): return fig +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 + ) + flat = axes.ravel() + for ax, key in zip(flat, ("rollout", "reference")): + side = r.payload.get(key, {}) + centers = np.asarray(side.get("centers", [])) + means = np.asarray(side.get("means", [])) + 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}" + ) + cum = cum + means[:, i] + if log_x: + ax.set_xscale("log") + ax.set_ylim(0, 1) + ax.set_title(_SERIES_LABELS[key], fontsize=8) + ax.set_xlabel(r.xlabel) + flat[0].set_ylabel("mean gate weight") + ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router") + return fig + + +def _render_router_share(r: Reduced, params: dict): + categories = r.payload["categories"] + n_experts = r.payload["n_experts"] + x = np.arange(len(categories)) + present = [k for k in ("rollout", "reference") if k in r.payload] + fig, axes = ps.new_figure( + "slide-16x9", + title=r.title, + params=params, + nrows=1, + ncols=len(present), + squeeze=False, + ) + flat = axes.ravel() + for ax, key in zip(flat, present): + side = r.payload[key] + shares = np.array([side[c] for c in categories]) # (n_cat, n_experts) + bottom = np.zeros(len(categories)) + for i in range(n_experts): + ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}") + bottom += shares[:, i] + ax.set_xticks(x) + ax.set_xticklabels(categories, rotation=45, ha="right") + ax.set_ylim(0, 1) + ax.set_title(_SERIES_LABELS[key], fontsize=8) + flat[0].set_ylabel("share of rows dispatched to expert") + ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router") + return fig + + +def _render_unavailable(r: Reduced, params: dict): + fig, ax = ps.new_figure("thesis-single", title=r.title, params=params) + ax.axis("off") + ax.text( + 0.5, + 0.5, + r.payload.get("note", "not available"), + ha="center", + va="center", + wrap=True, + fontsize=10, + transform=ax.transAxes, + ) + return fig + + _RENDERERS = { "overlay_hist": _render_overlay, "single_hist": _render_single, "grouped_hist": _render_grouped, "profile": _render_profile, "bar": _render_bar, + "router_gating": _render_router_gating, + "router_share": _render_router_share, + "unavailable": _render_unavailable, } diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py new file mode 100644 index 0000000..6dbc3f8 --- /dev/null +++ b/giant/analysis/router_gating.py @@ -0,0 +1,339 @@ +"""Router gating diagnostic: where a MoE checkpoint's decision boundaries sit. + +Unlike everything else in this package, this reduction needs a live PyTorch +model — soft expert gate weights aren't columns in a rollout/predict parquet, +they only exist by calling `Router.gate(cond_cont, cond_cat)` (see +`giant.model.network.Router`) against the checkpoint that produced the +rollout. That's a deliberate, narrow exception to the rest of the catalog's +"polars/numpy only" contract; it still runs fine as a `compute-one` HTCondor +job since torch is already installed there (the same env trains checkpoints). + +The routing axis is fixed to pre-step energy: every router type at least +indirectly depends on it (`EnergyRouter` reads it directly; `PdgRouter` and +`ProcessRouter` correlate with it through the physics), and it's the one axis +a reader can interpret without knowing the checkpoint's specific router +config. `x` is binned into equal-population (quantile) bins rather than +equal-width ones, since energy is heavy-tailed and equal-width bins would +leave the upper end almost empty. Mean gate weight per bin is stacked as +filled areas per expert — since `gate` rows are a partition of unity, the +stack always fills exactly to 1, and the crossover bands are the router's +soft decision boundaries (where two experts' means cross ~0.5). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import polars as pl + +from giant.analysis.grouping import pdg_label +from giant.analysis.reduced import Reduced + +if TYPE_CHECKING: + import torch + + from giant.data.transforms import Normalizer + +_SAMPLE_ROWS = 200_000 +_N_BINS = 40 +_TOP_K_PROCESS = 8 + +_COLS = ( + "pre_x", + "pre_y", + "pre_z", + "pre_E", + "pre_dx", + "pre_dy", + "pre_dz", + "layer_id", + "pdg", + "material", +) + + +@dataclass +class _RouterHandle: + router: "torch.nn.Module" + pdg_map: dict[int, int] + mat_map: dict[str, int] + cond_normalizer: "Normalizer" + conditioning: str + router_type: str + + +def load_router(checkpoint: str | Path) -> _RouterHandle | None: + """Load a checkpoint's Stage-1 router, or None if it isn't a MoE checkpoint.""" + import torch + + from giant.data.transforms import Normalizer + from giant.model.network import build_models + + ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) + model_cfg = ckpt.get("model_config") or {} + router_cfg = model_cfg.get("router") + if not router_cfg or not router_cfg.get("enabled"): + return None + + stage1, _ = build_models(model_cfg) + stage1.load_state_dict(ckpt["model"]) + stage1.eval() + + return _RouterHandle( + router=stage1.router, + pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()}, + mat_map={str(k): v for k, v in ckpt["mat_map"].items()}, + cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]), + conditioning=model_cfg.get("conditioning", "embedding"), + router_type=router_cfg["type"], + ) + + +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) + lf = lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold) + return lf.select(*_COLS, *extra_cols).collect(engine="streaming") + + +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 + training vocab can't be embedded — dropped here the same way + `giant.rollout`'s own known-pdg gate drops them at inference. The + returned df keeps every original column (filtered to the same rows), so + callers can key gate weights by any of them (energy, pdg, process, ...). + """ + import torch + + from giant.data.transforms import build_cond_features + + known = np.array( + [ + int(p) in handle.pdg_map and str(m) in handle.mat_map + for p, m in zip(df["pdg"].to_list(), df["material"].to_list()) + ] + ) + if not known.any(): + return df.clear(), np.zeros((0, handle.router.n_experts)) + 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_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()] + ), + "layer_id": df["layer_id"].to_numpy(), + "pdg": df["pdg"].to_numpy(), + "material": df["material"].to_numpy(), + } + cond_cont, cond_cat = build_cond_features( + data, + handle.pdg_map, + handle.mat_map, + cond_normalizer=handle.cond_normalizer, + conditioning=handle.conditioning, + ) + with torch.no_grad(): + gate = handle.router.gate( + torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long() + ).numpy() + return df, gate + + +def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict: + order = np.argsort(x) + x_sorted, g_sorted = x[order], gate[order] + edges = np.quantile(x_sorted, np.linspace(0, 1, n_bins + 1)) + edges[-1] = np.nextafter(edges[-1], np.inf) # include the max value + bin_idx = np.clip(np.digitize(x_sorted, edges[1:-1]), 0, n_bins - 1) + + n_experts = gate.shape[1] + centers = np.full(n_bins, np.nan) + means = np.full((n_bins, n_experts), np.nan) + for b in range(n_bins): + mask = bin_idx == b + if mask.any(): + centers[b] = x_sorted[mask].mean() + means[b] = g_sorted[mask].mean(axis=0) + valid = ~np.isnan(centers) + 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]]: + """Fraction of each category's rows hard-dispatched to each expert. + + Uses `Router.top1` (argmax), not the soft `gate` mean — grouped top-1 + dispatch is what `_route_forward` actually runs in eval mode (rollout, + predict), so this answers "which expert does a photon/Compton step + actually go through", not just its average soft weight. + """ + shares: dict[str, list[float]] = {} + for key in order: + mask = categories == key + total = int(mask.sum()) + if total == 0: + shares[str(key)] = [0.0] * n_experts + continue + counts = np.bincount(idx[mask], minlength=n_experts) + shares[str(key)] = (counts / total).tolist() + return shares + + +_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)", + "router_share_by_pdg": "Router expert share by particle species", + "router_share_by_process": "Router expert share by physics process", +} + + +def _unavailable(spec_id: str) -> Reduced: + return Reduced( + id=spec_id, + family="model", + kind="unavailable", + title=_TITLES[spec_id], + xlabel="n/a", + payload={"note": _NOTE_NOT_MOE}, + ) + + +def compute_router_gating( + checkpoint: str | Path | None, + r_phys: pl.LazyFrame, + t_phys: pl.LazyFrame, + seed: int = 0, +) -> Reduced: + """`Reduced` for the router-gating figure, or an explanatory note if n/a.""" + handle = load_router(checkpoint) if checkpoint else None + if handle is None: + return _unavailable("router_gating") + + sides: dict[str, dict] = {} + for name, lf in (("rollout", r_phys), ("reference", t_phys)): + 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": []} + ) + + return Reduced( + id="router_gating", + family="model", + kind="router_gating", + title=_TITLES["router_gating"], + xlabel="pre-step energy [MeV]", + payload={ + "router_type": handle.router_type, + "n_experts": handle.router.n_experts, + "log_x": True, + **sides, + }, + ) + + +def compute_router_share_by_pdg( + checkpoint: str | Path | None, + r_phys: pl.LazyFrame, + t_phys: pl.LazyFrame, + top_pdgs: list[int], + seed: int = 0, +) -> Reduced: + """Stacked-bar share of each particle species dispatched to each expert.""" + handle = load_router(checkpoint) if checkpoint else None + if handle is None: + return _unavailable("router_share_by_pdg") + + labels = [pdg_label(p) for p in top_pdgs] + sides: dict[str, dict] = {} + for name, lf in (("rollout", r_phys), ("reference", t_phys)): + df = _subsample(lf, _SAMPLE_ROWS, seed) + 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 + ) + 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)} + + return Reduced( + id="router_share_by_pdg", + family="model", + kind="router_share", + title=_TITLES["router_share_by_pdg"], + xlabel="particle species", + payload={ + "router_type": handle.router_type, + "n_experts": handle.router.n_experts, + "categories": labels, + **sides, + }, + ) + + +def compute_router_share_by_process( + checkpoint: str | Path | None, + t_phys: pl.LazyFrame, + seed: int = 0, + top_k: int = _TOP_K_PROCESS, +) -> Reduced: + """Stacked-bar share of each physics process dispatched to each expert. + + Reference-only: ``process`` is the true post-step physics process — a + label the rollout side has no equivalent of (see + `giant.model.network.ProcessRouter`, which predicts it from pre-step + conditioning alone, never observes it at eval time). This plot instead + checks *after the fact*, on real data, how well the router's conditioning + -based dispatch lines up with the true process. + """ + handle = load_router(checkpoint) if checkpoint else None + if handle is None: + return _unavailable("router_share_by_process") + + df = _subsample(t_phys, _SAMPLE_ROWS, seed, extra_cols=("process",)) + df, gate = _gate_for_df(handle, df) + if len(df): + 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 + ) + else: + order, shares = [], {} + + return Reduced( + id="router_share_by_process", + family="model", + kind="router_share", + title=_TITLES["router_share_by_process"], + xlabel="physics process", + payload={ + "router_type": handle.router_type, + "n_experts": handle.router.n_experts, + "categories": order, + "reference": {p: shares[p] for p in order}, + }, + ) diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 02d449b..90447d3 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -44,6 +44,9 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle): "profile", "bar", "single_hist", + "router_gating", + "router_share", + "unavailable", } assert r.title and r.xlabel _validate_payload(r) @@ -67,3 +70,14 @@ def _validate_payload(r) -> None: assert len(p[k]) == n elif r.kind == "bar": assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"]) + elif r.kind == "unavailable": + assert p["note"] + elif r.kind == "router_gating": + for side in ("rollout", "reference"): + if side in p: + assert len(p[side]["centers"]) == len(p[side]["means"]) + elif r.kind == "router_share": + for cat in p["categories"]: + for side in ("rollout", "reference"): + if side in p: + assert cat in p[side] diff --git a/tests/test_router_gating.py b/tests/test_router_gating.py new file mode 100644 index 0000000..51123d3 --- /dev/null +++ b/tests/test_router_gating.py @@ -0,0 +1,126 @@ +"""Tests for the MoE router-gating diagnostic (giant.analysis.router_gating).""" + +from __future__ import annotations + +import numpy as np +import polars as pl +import torch + +from giant.analysis.router_gating import ( + compute_router_gating, + compute_router_share_by_pdg, + compute_router_share_by_process, +) +from giant.data.transforms import Normalizer +from giant.model.network import build_models + +_PDG_MAP = {11: 0, 22: 1} +_MAT_MAP = {"G4_PbWO4": 0, "G4_Pb": 1} + + +def _model_cfg() -> dict: + return { + "router": { + "enabled": True, + "type": "energy", + "n_experts": 2, + "temperature": 0.5, + "learn_centers": True, + "energy_idx": 3, + }, + "pdg_vocab": len(_PDG_MAP), + "mat_vocab": len(_MAT_MAP), + "conditioning": "embedding", + } + + +def _write_checkpoint(tmp_path) -> str: + cfg = _model_cfg() + stage1, _ = build_models(cfg) + norm = Normalizer() + norm.mean = np.zeros(15, dtype=np.float32) + norm.std = np.ones(15, dtype=np.float32) + ckpt = { + "model_config": cfg, + "model": stage1.state_dict(), + "pdg_map": _PDG_MAP, + "mat_map": _MAT_MAP, + "normalizer": {"cond": norm.to_dict()}, + } + path = tmp_path / "ckpt.pt" + torch.save(ckpt, path) + return str(path) + + +def _steps_frame(process: bool = False) -> pl.LazyFrame: + n = 40 + rng = np.random.default_rng(0) + pre_e = np.concatenate([rng.uniform(1, 10, n // 2), rng.uniform(100, 1000, n // 2)]) + pdg = np.where(np.arange(n) % 2 == 0, 11, 22) + material = np.where(np.arange(n) % 3 == 0, "G4_Pb", "G4_PbWO4") + data = { + "event_id": np.arange(n), + "pdg": pdg, + "pre_x": np.zeros(n), + "pre_y": np.zeros(n), + "pre_z": np.zeros(n), + "pre_E": pre_e, + "pre_dx": np.zeros(n), + "pre_dy": np.zeros(n), + "pre_dz": np.ones(n), + "post_x": np.zeros(n), + "post_y": np.zeros(n), + "post_z": np.ones(n), + "post_E": pre_e * 0.5, + "post_dx": np.zeros(n), + "post_dy": np.zeros(n), + "post_dz": np.ones(n), + "edep": pre_e * 0.5, + "step_length": np.ones(n), + "material": material, + "layer_id": np.zeros(n, dtype=np.int64), + } + if process: + data["process"] = np.where(pdg == 11, "eIoni", "compt") + return pl.DataFrame(data).lazy() + + +def test_compute_router_gating_shapes(tmp_path): + checkpoint = _write_checkpoint(tmp_path) + lf = _steps_frame() + r = compute_router_gating(checkpoint, lf, lf) + assert r.kind == "router_gating" + assert r.payload["n_experts"] == 2 + for side in ("rollout", "reference"): + means = r.payload[side]["means"] + assert means, f"{side} produced no bins" + assert all(abs(sum(row) - 1.0) < 1e-5 for row in means) + + +def test_compute_router_gating_missing_checkpoint_is_unavailable(): + lf = _steps_frame() + r = compute_router_gating(None, lf, lf) + assert r.kind == "unavailable" + assert "note" in r.payload + assert r.title + + +def test_compute_router_share_by_pdg(tmp_path): + checkpoint = _write_checkpoint(tmp_path) + lf = _steps_frame() + r = compute_router_share_by_pdg(checkpoint, lf, lf, top_pdgs=[11, 22]) + assert r.kind == "router_share" + for side in ("rollout", "reference"): + assert set(r.payload[side]) == {"e-", "gamma"} + for shares in r.payload[side].values(): + assert abs(sum(shares) - 1.0) < 1e-5 + + +def test_compute_router_share_by_process(tmp_path): + checkpoint = _write_checkpoint(tmp_path) + lf = _steps_frame(process=True) + r = compute_router_share_by_process(checkpoint, lf) + assert r.kind == "router_share" + assert set(r.payload["categories"]) <= {"eIoni", "compt"} + for shares in r.payload["reference"].values(): + assert abs(sum(shares) - 1.0) < 1e-5