"""Unit tests for v10 (parked Universal Causal-Hyperbolic) modules. Covers: causal SCM, geometric prior (shared with v9b), counterfactual decoder (shared with v9b), hyperbolic conformal, integrated v10 model. Run with: python tests/test_v10_components.py or via pytest: pytest tests/test_v10_components.py v10 is parked at src/research/_v10_universal_hyperbolic/. Active research direction is v9b (Normative JEPA + Conformal); see proposals/v9b_*.md. """ from __future__ import annotations import sys from pathlib import Path import torch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from src.research._v10_universal_hyperbolic.causal_scm import ( # noqa: E402 CausalSCMHead, CausalSplitHead, CausalRecompose, LearnableDAGAdjacency, orthogonality_loss, ) from src.research.geometric_prior import ( # noqa: E402 GeometricPriorConditioning, synthetic_brain_sdf_template, make_coord_grid, SIRENImplicitSDF, ) from src.research.counterfactual_decoder import ( # noqa: E402 CounterfactualHealthyDecoder, tumor_residual, ) from src.research._v10_universal_hyperbolic.hyperbolic_conformal import ( # noqa: E402 HyperbolicCalibSample, HyperbolicConformalCalibrator, voxelwise_hyperbolic_anomaly_map, weighted_quantile_tibshirani, ) from src.research._v10_universal_hyperbolic.hyperbolic import expmap0 # noqa: E402 # ---------------------------------------------------------------------- # Causal SCM # ---------------------------------------------------------------------- def test_causal_split_dims(): """SplitHead produces 3 streams with correct dims.""" head = CausalSplitHead(in_dim=256, anatomy_dim=128, tumor_dim=64, scanner_dim=32) z = torch.randn(8, 256) z_a, z_t, z_s = head(z) assert z_a.shape == (8, 128) assert z_t.shape == (8, 64) assert z_s.shape == (8, 32) def test_orthogonality_loss(): """Orthogonality loss (cross-correlation Frobenius norm) is small for independent vectors, larger for correlated, and works on different dims. For identical vectors of dim D: diagonal of normalized cross-cov is 1, off-diagonal small, mean of squared ~ D/D^2 = 1/D. For independent random vectors of dim D, samples N: all entries small (concentration of measure), mean of squared ~ 1/N. So parallel > orthogonal but the margin shrinks with D. """ n = 100 z1 = torch.randn(n, 32) # Orthogonal: independent random z2_ortho = torch.randn(n, 32) loss_ortho = orthogonality_loss(z1, z2_ortho) # Parallel: identical z2_parallel = z1.clone() loss_parallel = orthogonality_loss(z1, z2_parallel) assert loss_parallel > loss_ortho, \ f"parallel loss {loss_parallel.item():.4f} should exceed ortho {loss_ortho.item():.4f}" # Cross-correlation magnitude for identical 32-dim vectors: ~1/32 = 0.031. assert loss_parallel.item() > 0.02, \ f"parallel cross-correlation should be at least 1/D ~ 0.03, got {loss_parallel.item():.4f}" # Test mixed-dim case (the bug we fixed) z3 = torch.randn(n, 16) loss_mixed = orthogonality_loss(z1, z3) assert torch.isfinite(loss_mixed).all() and loss_mixed.item() >= 0, \ f"mixed-dim loss should be finite + non-negative, got {loss_mixed.item()}" assert loss_mixed.item() < 1.0, \ f"mixed-dim independent loss should be near 0, got {loss_mixed.item()}" def test_dag_acyclicity(): """NOTEARS h(A) is 0 for a DAG, > 0 for a cyclic adjacency.""" dag = LearnableDAGAdjacency() # Set A to be a clean DAG: anatomy -> tumor -> scanner only. A = torch.zeros(3, 3) A[0, 1] = 0.5 # anatomy -> tumor A[1, 2] = 0.5 # tumor -> scanner dag.raw.data = A h = dag.h().item() assert abs(h) < 0.5, f"clean DAG should give small h, got {h}" # Now inject a cycle: scanner -> anatomy. A_cyclic = A.clone() A_cyclic[2, 0] = 0.5 dag.raw.data = A_cyclic h_cyc = dag.h().item() assert h_cyc > h, f"cyclic A should give larger h ({h_cyc}) than acyclic ({h})" def test_scm_head_full_forward(): """Full SCM head produces recomposed latent + all aux losses.""" scm = CausalSCMHead(in_dim=256, anatomy_dim=128, tumor_dim=64, scanner_dim=32, decoder_in_dim=256) z = torch.randn(8, 256) recomposed, aux = scm(z) assert recomposed.shape == (8, 256) for k in ("ortho_at", "ortho_as", "ortho_ts", "dag", "dag_forbidden", "z_anatomy", "z_tumor", "z_scanner"): assert k in aux, f"missing aux key: {k}" # Counterfactual: z_tumor should be zeroed recomposed_cf, _ = scm(z, counterfactual_healthy=True) # Different recomposed because z_tumor is zeroed assert not torch.allclose(recomposed, recomposed_cf) # ---------------------------------------------------------------------- # Geometric prior # ---------------------------------------------------------------------- def test_synthetic_brain_sdf(): """SDF template has correct sign convention.""" sdf = synthetic_brain_sdf_template(size=64) # Center should be inside (negative SDF) center = sdf[32, 32].item() # Corner should be outside (positive SDF) corner = sdf[0, 0].item() assert center < 0, f"center SDF should be negative (inside), got {center}" assert corner > 0, f"corner SDF should be positive (outside), got {corner}" def test_geometric_prior_concat_mode(): """Concat mode adds 1 channel.""" prior = GeometricPriorConditioning(image_size=64, mode="concat") x = torch.randn(2, 3, 64, 64) out = prior(x) assert out.shape == (2, 4, 64, 64), f"expected (2, 4, 64, 64), got {out.shape}" def test_geometric_prior_blend_mode(): """Blend mode preserves channel count.""" prior = GeometricPriorConditioning(image_size=64, mode="blend", refine=False) x = torch.randn(2, 3, 64, 64) out = prior(x) assert out.shape == (2, 3, 64, 64) def test_siren_implicit_sdf(): """SIREN INR output matches expected dimensions.""" inr = SIRENImplicitSDF(hidden_dim=64, n_layers=3) coords = make_coord_grid(16) out = inr(coords) assert out.shape == (16, 16, 1) # ---------------------------------------------------------------------- # Counterfactual decoder # ---------------------------------------------------------------------- def test_counterfactual_decoder_shape(): """Decoder produces image-shaped output.""" dec = CounterfactualHealthyDecoder(latent_dim=192, image_size=64) x = torch.randn(2, 3, 64, 64) z = torch.randn(2, 192) out = dec(x, z) assert out.shape == (2, 3, 64, 64) assert (out >= -1).all() and (out <= 1).all(), "tanh output should be in [-1, 1]" def test_counterfactual_reconstruction_loss(): """Reconstruction loss is 0 when x_healthy == x_input.""" dec = CounterfactualHealthyDecoder(latent_dim=192, image_size=64) x = torch.randn(2, 3, 64, 64) mask = torch.zeros(2, 1, 64, 64) # healthy scan loss = dec.reconstruction_loss(x, x, mask) assert loss.item() < 1e-5, f"loss should be ~0 when x_healthy == x, got {loss}" def test_tumor_residual(): """Residual is high where input differs from counterfactual.""" x = torch.zeros(1, 3, 64, 64) x_cf = torch.zeros(1, 3, 64, 64) x[0, :, 20:40, 20:40] = 1.0 # "tumor" in center residual = tumor_residual(x, x_cf) assert residual[0, 0, 30, 30].item() > 0.5, "residual should be high in tumor region" assert residual[0, 0, 5, 5].item() < 0.1, "residual should be low outside tumor" # ---------------------------------------------------------------------- # Hyperbolic conformal # ---------------------------------------------------------------------- def test_weighted_quantile_uniform(): """Uniform weights should match standard quantile.""" values = torch.linspace(0, 1, 100).numpy() weights = torch.ones(100).numpy() q = weighted_quantile_tibshirani(values, weights, 0.9) # With inflation, should be near 0.91 (slightly above 0.9 by Tibshirani correction) assert 0.85 < q < 0.95, f"expected quantile ~0.9, got {q}" def test_hyperbolic_calibrator(): """Calibrator empirical coverage matches target (1-alpha).""" import numpy as np torch.manual_seed(0) np.random.seed(0) # Synthetic calibration: predict z_pred = z_true + small noise on Poincare ball samples = [] for _ in range(200): z_true_eu = torch.randn(16) * 0.3 z_true = expmap0(z_true_eu, c=1.0) # Predicted = true + noise (in tangent space) z_pred_eu = z_true_eu + torch.randn(16) * 0.1 z_pred = expmap0(z_pred_eu, c=1.0) samples.append(HyperbolicCalibSample(z_pred=z_pred, z_true=z_true)) cal = HyperbolicConformalCalibrator(alpha=0.10, curvature_c=1.0) report = cal.calibrate(samples) assert 0.85 <= report.empirical_coverage_on_calib <= 1.0, \ f"coverage {report.empirical_coverage_on_calib} should be >= 0.85 (target 0.90)" assert cal.q > 0 assert cal.q < 5 # reasonable for unit Poincare ball def test_hyperbolic_calibrator_predict(): """Calibrator predict produces well-formed output dict.""" torch.manual_seed(1) samples = [] for _ in range(50): z_true = expmap0(torch.randn(8) * 0.3, c=1.0) z_pred = expmap0(torch.randn(8) * 0.3, c=1.0) samples.append(HyperbolicCalibSample(z_pred=z_pred, z_true=z_true)) cal = HyperbolicConformalCalibrator(alpha=0.10, curvature_c=1.0) cal.calibrate(samples) z_pred_test = expmap0(torch.randn(8) * 0.3, c=1.0) z_true_test = expmap0(torch.randn(8) * 0.3, c=1.0) out = cal.predict(z_pred_test, z_true_test) for k in ("score", "q", "in_prediction_set", "anomaly_certified"): assert k in out assert isinstance(out["in_prediction_set"], bool) def test_voxelwise_anomaly_map(): """Voxelwise wrapper produces correctly-shaped boolean map.""" torch.manual_seed(2) samples = [] for _ in range(50): z_true = expmap0(torch.randn(4) * 0.3, c=1.0) z_pred = expmap0(torch.randn(4) * 0.3, c=1.0) samples.append(HyperbolicCalibSample(z_pred=z_pred, z_true=z_true)) cal = HyperbolicConformalCalibrator(alpha=0.10, curvature_c=1.0) cal.calibrate(samples) z_pred_vox = expmap0(torch.randn(2, 4, 16, 16) * 0.3, c=1.0) z_tpl_vox = expmap0(torch.randn(2, 4, 16, 16) * 0.3, c=1.0) amap = voxelwise_hyperbolic_anomaly_map(z_pred_vox, z_tpl_vox, cal) assert amap.shape == (2, 1, 16, 16) assert amap.dtype == torch.bool # ---------------------------------------------------------------------- # V10 integrated model # ---------------------------------------------------------------------- def test_v10_model_forward_smoke(): """V10Model end-to-end forward produces all expected outputs.""" try: import segmentation_models_pytorch as smp # noqa except ImportError: print("SMP not installed, skipping integrated model test") return from src.research._v10_universal_hyperbolic.v10_model import V10Model # Small image size to keep test fast; ConvNeXt-Tiny still loads. model = V10Model( image_size=128, latent_dim=128, anatomy_dim=64, tumor_dim=32, scanner_dim=16, use_counterfactual=True, use_geometric_prior=True, ).eval() x = torch.randn(2, 3, 128, 128) with torch.no_grad(): out = model(x, return_counterfactual=True) expected_keys = ["mask_logits", "z_euclidean", "z_hyperbolic", "z_tangent", "z_anatomy", "z_tumor", "z_scanner", "x_counterfactual", "tumor_residual", "hyperbolic_curvature", "aux_losses"] for k in expected_keys: assert k in out, f"missing key {k}" assert out["mask_logits"].shape == (2, 1, 128, 128) assert out["x_counterfactual"].shape == (2, 3, 128, 128) assert out["tumor_residual"].shape == (2, 1, 128, 128) assert out["z_anatomy"].shape == (2, 64) assert out["z_tumor"].shape == (2, 32) assert out["z_scanner"].shape == (2, 16) # Hyperbolic latent must be inside the Poincare ball z_h_norm = out["z_hyperbolic"].norm(dim=-1) assert (z_h_norm < 1.0).all(), \ f"hyperbolic latent should be in unit ball, max norm {z_h_norm.max()}" if __name__ == "__main__": print("=== causal SCM ===") test_causal_split_dims() test_orthogonality_loss() test_dag_acyclicity() test_scm_head_full_forward() print("=== geometric prior ===") test_synthetic_brain_sdf() test_geometric_prior_concat_mode() test_geometric_prior_blend_mode() test_siren_implicit_sdf() print("=== counterfactual decoder ===") test_counterfactual_decoder_shape() test_counterfactual_reconstruction_loss() test_tumor_residual() print("=== hyperbolic conformal ===") test_weighted_quantile_uniform() test_hyperbolic_calibrator() test_hyperbolic_calibrator_predict() test_voxelwise_anomaly_map() print("=== v10 integrated model ===") test_v10_model_forward_smoke() print("\nAll v10 component tests passed.")