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:
+14
-3
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user