Bump ruff line-length to 120 and reformat

Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.
2026-08-12 13:33:09 +02:00
parent 9ce7b32324
commit 55332db67a
66 changed files with 757 additions and 2413 deletions
@@ -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
@@ -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"


@@ -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
@@ -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()

@@ -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("/")
@@ -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
@@ -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 == ""
@@ -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:
@@ -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

@@ -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()
@@ -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):
@@ -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
@@ -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})")


@@ -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)
@@ -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}


@@ -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)
@@ -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)
@@ -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"}
@@ -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())
@@ -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)

@@ -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]]
@@ -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 ────────────────────────────────────────────────────
@@ -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]]
@@ -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):
@@ -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)


@@ -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:
@@ -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)
@@ -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())