lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Bump ruff line-length to 120 and reformat
CI / Lint (ruff check) (push) Successful in 31s
Details
CI / Format (ruff format) (push) Successful in 32s
Details
CI / Sync project version with tag (push) Has been skipped
Details
CI / Lint (ruff check) (pull_request) Successful in 26s
Details
CI / Type check (ty) (push) Successful in 29s
Details
CI / Format (ruff format) (pull_request) Successful in 33s
Details
CI / Sync project version with tag (pull_request) Has been skipped
Details
CI / Type check (ty) (pull_request) Successful in 35s
Details
CI / Tests (pull_request) Successful in 3m47s
Details
CI / Tests (push) Successful in 3m55s
Details

Browse Source 
Rejoins lines that only wrapped because they exceeded the old 88-char
limit; ruff check and the full test suite (725 passed) are unaffected.
This commit is contained in:
lars
2026-08-12 13:33:09 +02:00
parent 9ce7b32324
commit 55332db67a
66 changed files with 757 additions and 2413 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_phase2.py
+20 -60
View File
@@ -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)