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