Apply ruff format across the codebase

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

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

tests/test_phase2.py

@@ -1,4 +1,5 @@
 """Tests for Phase 2: secondary particle prediction."""
+
 import numpy as np
 import pytest
 import torch
@@ -11,6 +12,7 @@ from giant.sample import sample_secondaries, snap_type_to_pdg_idx
 
 # ── helpers ──────────────────────────────────────────────────────────────────
 
+
 def _stage1(pdg=3, mat=2):
     return DenoisingMLP(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2)
 
@@ -29,6 +31,7 @@ def _cond(B=8, pdg=3, mat=2):
 
 # ── DenoisingMLP Phase-2 additions ───────────────────────────────────────────
 
+
 def test_predict_n_sec_shape():
     B = 8
     model = _stage1()
@@ -53,6 +56,7 @@ def test_pdg_embedding_weight_shape():
 
 # ── SecondaryDecoder ──────────────────────────────────────────────────────────
 
+
 def test_sec_decoder_output_shape():
     B = 8
     decoder = _sec_decoder()
@@ -89,6 +93,7 @@ def test_sec_decoder_gradients():
 
 # ── masked flow matching loss ─────────────────────────────────────────────────
 
+
 def test_flow_matching_loss_secondary_scalar():
     B, pdg, mat = 8, 3, 2
     decoder = _sec_decoder(pdg, mat)
@@ -96,7 +101,9 @@ def test_flow_matching_loss_secondary_scalar():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
-    loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
+    loss = flow_matching_loss_secondary(
+        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
+    )
     assert loss.shape == ()
     assert loss.item() >= 0.0
 
@@ -109,7 +116,9 @@ def test_flow_matching_loss_secondary_mask_zeros_padding():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
-    loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
+    loss = flow_matching_loss_secondary(
+        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
+    )
     assert loss.item() == pytest.approx(0.0, abs=1e-6)
 
 
@@ -128,6 +137,7 @@ def test_flow_matching_loss_secondary_has_grad():
 
 # ── sampling ──────────────────────────────────────────────────────────────────
 
+
 def test_sample_secondaries_shapes():
     B, pdg, mat = 6, 3, 2
     decoder = _sec_decoder(pdg, mat)
@@ -169,6 +179,7 @@ def test_snap_type_to_pdg_idx_shape():
 
 # ── encode_secondaries round-trip ─────────────────────────────────────────────
 
+
 def test_encode_secondaries_energy_conservation():
     """Decoded stick-breaking fractions must sum to ≈ e_sec."""
     from giant.data.transforms import encode_secondaries
@@ -217,6 +228,6 @@ def test_encode_secondaries_direction_encoding():
     sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
 
     # dir columns are sec_cont[:, :, 1:4]
-    local_dirs = sec_cont[:, :2, 1:4]   # (N, 2, 3) — valid slots only
+    local_dirs = sec_cont[:, :2, 1:4]  # (N, 2, 3) — valid slots only
     norms_out = np.linalg.norm(local_dirs, axis=-1)
     np.testing.assert_allclose(norms_out, 1.0, atol=1e-5)

tests/test_rollout.py

@@ -56,16 +56,31 @@ def _seeds(n=6):
     }
 
 
-def _run(escape_threshold=1e9, energy_cutoff=1.0, max_steps=30,
-         max_tracks_per_event=300, seeds=None):
+def _run(
+    escape_threshold=1e9,
+    energy_cutoff=1.0,
+    max_steps=30,
+    max_tracks_per_event=300,
+    seeds=None,
+):
     torch.manual_seed(0)
     np.random.seed(0)
     s1, s2 = _models()
     cond, tgt = _norms()
     return rollout(
-        s1, s2, _oracle(), seeds or _seeds(), cond, tgt, PDG_MAP, MAT_MAP,
-        energy_cutoff=energy_cutoff, max_steps=max_steps, steps=4,
-        batch_size=128, max_tracks_per_event=max_tracks_per_event,
+        s1,
+        s2,
+        _oracle(),
+        seeds or _seeds(),
+        cond,
+        tgt,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=energy_cutoff,
+        max_steps=max_steps,
+        steps=4,
+        batch_size=128,
+        max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
     )