ctt progress 2026-07-02 workspace/tests/test_ctt.py

workspace/tests/test_ctt.py

@@ -3,7 +3,7 @@ import pytest
 torch = pytest.importorskip("torch")
 
 from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer
-from cil.models.ctt import chamfer_to_target_set, negative_boundary_loss
+from cil.models.ctt import chamfer_to_target_set, diversity_loss, negative_boundary_loss
 
 
 def test_ctt_variants_preserve_tangent_shape():
@@ -31,6 +31,24 @@ def test_chart_encoder_and_losses_are_finite():
 
     assert torch.isfinite(chamfer_to_target_set(predicted, positives))
     assert negative_boundary_loss(predicted, negatives, margin=0.2).item() == pytest.approx(0.0)
+    assert diversity_loss(predicted).item() >= 0.0
+
+
+def test_gated_residual_matches_documented_formula():
+    model = CausalTangentTransport(
+        CTTConfig(chart_feature_dim=10, chart_dim=2, tangent_dim=2, hidden_dim=4, variant="gated_residual")
+    )
+    for parameter in model.delta.parameters():
+        parameter.data.zero_()
+    for parameter in model.gate.parameters():
+        parameter.data.zero_()
+    model.delta[-1].bias.data[:] = torch.tensor([4.0, -2.0])
+    model.gate[2].bias.data[:] = torch.tensor([0.0, 0.0])
+
+    xi = torch.tensor([[2.0, 6.0]])
+    output = model(torch.zeros(1, 2), torch.zeros(1, 2), xi)
+
+    assert torch.allclose(output, torch.tensor([[3.0, 2.0]]), atol=1.0e-6)
 
 
 def test_tangent_normalizer_round_trips():