Encode edep/secondary/post energy as a conservation-constrained simplex

Replaces the independent log_delta_e/log_edep targets with 2 additive-log-ratio
coordinates over the deposit/secondary/post-energy simplex (fractions of pre_E
summing to 1), so edep + e_sec + post_E == pre_E holds by construction after
decoding (softmax) rather than being learned approximately. Requires e_sec
(secondary energy) as a new conditioning input and a steps_to_parquet.py pass
to derive it from child track first-step energies.
This commit is contained in:
2026-06-25 16:01:13 +02:00
parent 64c6bd1cef
commit 8475199609
13 changed files with 415 additions and 83 deletions
+3 -2
View File
@@ -1,4 +1,5 @@
import torch
from giant.constants import COND_DIM
from giant.model.network import DenoisingMLP, SinusoidalEmbedding
@@ -19,7 +20,7 @@ def test_denoising_mlp_output_shape():
model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
x_t = torch.randn(B, 9)
t = torch.rand(B)
cond_cont = torch.randn(B, 9)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[
torch.randint(0, 5, (B,)),
@@ -36,7 +37,7 @@ def test_denoising_mlp_gradients_flow():
model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
x_t = torch.randn(B, 9)
t = torch.rand(B)
cond_cont = torch.randn(B, 9)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
loss = model(x_t, t, cond_cont, cond_cat).sum()
loss.backward()