ctt train-calibration hygiene 2026-07-03T16:31:19Z: tests/test_ctt.py

tests/test_ctt.py

@@ -0,0 +1,92 @@
+import pytest
+import numpy as np
+from pathlib import Path
+
+torch = pytest.importorskip("torch")
+
+from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer
+from cil.models.ctt import chamfer_to_target_set, diversity_loss, negative_boundary_loss
+from scripts.eval_ctt_generated_rollout import ChartItem, _source_pool_for_target
+
+
+def test_ctt_variants_preserve_tangent_shape():
+    z_source = torch.randn(3, 8)
+    z_target = torch.randn(3, 8)
+    xi_source = torch.randn(3, 5)
+
+    for variant in ("residual", "gated_residual"):
+        model = CausalTangentTransport(
+            CTTConfig(chart_feature_dim=10, chart_dim=8, tangent_dim=5, hidden_dim=16, variant=variant)
+        )
+        output = model(z_source, z_target, xi_source)
+        assert output.shape == xi_source.shape
+        assert torch.isfinite(output).all()
+
+
+def test_chart_encoder_and_losses_are_finite():
+    encoder = ChartEncoder(input_dim=6, hidden_dim=12, output_dim=8)
+    z = encoder(torch.randn(4, 6))
+    assert z.shape == (4, 8)
+
+    predicted = torch.randn(4, 5)
+    positives = predicted + 0.1
+    negatives = predicted + 10.0
+
+    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():
+    values = torch.randn(8, 5)
+    normalizer = TangentNormalizer.fit(values)
+    restored = normalizer.inverse_transform(normalizer.transform(values))
+    assert torch.allclose(values, restored, atol=1.0e-5)
+
+
+def test_rollout_source_pool_excludes_target_chart_and_state_hash():
+    def chart(chart_id: str, state_hash: str) -> ChartItem:
+        return ChartItem(
+            chart_id=chart_id,
+            task_id="task",
+            seed="0",
+            state_hash=state_hash,
+            instruction="",
+            source_dataset=Path("."),
+            base_action=np.zeros((2, 7), dtype=np.float32),
+            feature=np.zeros(4, dtype=np.float32),
+            positive_tangents=np.zeros((1, 21), dtype=np.float32),
+            negative_tangents=np.zeros((0, 21), dtype=np.float32),
+            hidden_utilities=[],
+            hidden_candidate_types=[],
+            stored_base_utility=None,
+        )
+
+    target = chart("chart_a", "state_a")
+    sources = [chart("chart_a", "state_a"), chart("chart_b", "state_a"), chart("chart_c", "state_c")]
+
+    pool = _source_pool_for_target(
+        target,
+        task_pool=sources,
+        source_charts=sources,
+        exclude_self_source=True,
+    )
+
+    assert [item.chart_id for item in pool] == ["chart_c"]