diff options
Diffstat (limited to 'tests/test_core.py')
| -rw-r--r-- | tests/test_core.py | 236 |
1 files changed, 236 insertions, 0 deletions
diff --git a/tests/test_core.py b/tests/test_core.py new file mode 100644 index 0000000..50d3aad --- /dev/null +++ b/tests/test_core.py @@ -0,0 +1,236 @@ +import numpy as np +import torch + +from worldalign.common import retrieval_metrics, sliced_wasserstein +from worldalign.energy import ( + log_sinkhorn, + relation_field_energy, + standardized_relation, +) +from worldalign.gw import gw_pseudo_targets +from worldalign.models import Bridge, PrefixAdapter +from worldalign.vg_diagnose import bundle_signature + + +def test_shapes_and_retrieval(): + bridge = Bridge(8, 12, hidden_dim=16) + prefix = PrefixAdapter(12, 20, prefix_length=4, hidden_dim=16) + x = torch.randn(5, 8) + y = bridge(x) + assert y.shape == (5, 12) + assert prefix(y).shape == (5, 4, 20) + metrics = retrieval_metrics(y, y) + assert metrics["i2t_r@1"] == 1.0 + + +def test_sliced_wasserstein_identity(): + x = torch.randn(32, 16) + assert sliced_wasserstein(x, x).item() < 1e-10 + + +def test_gw_recovers_structure_without_pairs(): + rng = np.random.default_rng(7) + z = rng.normal(size=(600, 6)).astype(np.float32) + q, _ = np.linalg.qr(rng.normal(size=(6, 6))) + x = torch.from_numpy(z) + y = torch.from_numpy((z @ q).astype(np.float32))[torch.randperm(len(z))] + result = gw_pseudo_targets(x, y, clusters=12, seed=3, max_iter=50) + assert result["coupling"].shape == (12, 12) + assert np.isfinite(result["gw_distance"]) + + +def test_bundle_signature_is_view_permutation_and_rotation_invariant(): + torch.manual_seed(5) + x = torch.randn(7, 8, 12) + q, _ = torch.linalg.qr(torch.randn(12, 12)) + permuted = (x @ q)[:, torch.randperm(8)] + first = bundle_signature(x) + second = bundle_signature(permuted) + assert torch.allclose(first, second, atol=2e-5) + + +def test_log_sinkhorn_has_unit_marginals(): + logits = torch.randn(9, 9) + coupling = log_sinkhorn(logits, temperature=0.7, iterations=30) + assert torch.allclose(coupling.sum(0), torch.ones(9), atol=1e-4) + assert torch.allclose(coupling.sum(1), torch.ones(9), atol=1e-4) + + +def test_transposition_delta_matches_brute_force(): + from worldalign.manifold_gate import ( + all_transposition_delta_mse, + build_channels, + permuted, + relation_mse, + ) + + generator = torch.Generator().manual_seed(11) + visual = torch.randn(24, 8, generator=generator)[:, None, :] + text = torch.randn(24, 10, generator=generator)[:, None, :] + text_channels, _ = build_channels(text, bundle=False) + visual_channels, _ = build_channels(visual, bundle=False) + delta = all_transposition_delta_mse(text_channels[0], visual_channels[0]) + base = relation_mse(text_channels[0], visual_channels[0]) + for p, q in [(0, 1), (3, 17), (5, 23), (10, 11)]: + permutation = torch.arange(24) + permutation[[p, q]] = permutation[[q, p]] + brute = ( + relation_mse(permuted(text_channels[0], permutation), visual_channels[0]) + - base + ) + assert torch.allclose(delta[p, q], brute, atol=1e-10) + assert torch.allclose(delta[q, p], brute, atol=1e-10) + + +def test_gate_energy_matches_energy_module(): + from worldalign.manifold_gate import assignment_energy, build_channels + + generator = torch.Generator().manual_seed(3) + visual = torch.randn(16, 6, generator=generator) + text = torch.randn(16, 9, generator=generator) + text_channels, text_relation = build_channels(text[:, None, :], bundle=False) + visual_channels, visual_relation = build_channels(visual[:, None, :], bundle=False) + permutation = torch.randperm(16, generator=generator) + gate = assignment_energy( + text_channels, visual_channels, text_relation, visual_relation, permutation + ) + reference_relation, reference_standardized = standardized_relation(visual) + mse, kl = relation_field_energy( + reference_relation, reference_standardized, text[permutation] + ) + assert abs(gate["mse"] - float(mse)) < 1e-5 + assert abs(gate["conditional_kl"] - float(kl)) < 1e-4 + + +def test_k_derangement_moves_exactly_k(): + from worldalign.manifold_gate import k_derangement + + generator = torch.Generator().manual_seed(9) + for k in (2, 4, 16): + permutation = k_derangement(64, k, generator) + moved = (permutation != torch.arange(64)).sum() + assert int(moved) == k + + +def test_bundle_mean_channel_matches_view_mean_relation(): + from worldalign.manifold_gate import view_bundle_channels + + torch.manual_seed(4) + views = torch.nn.functional.normalize(torch.randn(6, 5, 8), dim=-1) + channels = view_bundle_channels(views) + assert channels.shape == (4, 6, 6) + means = views.double().mean(1) + assert torch.allclose(channels[0], means @ means.T, atol=1e-10) + + +def test_batched_transposition_delta_matches_single(): + from worldalign.manifold_gate import all_transposition_delta_mse + from worldalign.view_gate import batched_transposition_delta, standardized_field + + torch.manual_seed(2) + views_a = torch.nn.functional.normalize(torch.randn(3, 6, 5), dim=-1) + views_b = torch.nn.functional.normalize(torch.randn(3, 6, 7), dim=-1) + fields_a = standardized_field(views_a) + fields_b = standardized_field(views_b) + batched = batched_transposition_delta(fields_a, fields_b) + for item in range(3): + single = all_transposition_delta_mse(fields_a[item], fields_b[item]) + assert torch.allclose(batched[item], single, atol=1e-10) + + +def test_batched_descent_reaches_local_minimum(): + from worldalign.view_gate import ( + batched_descent, + batched_transposition_delta, + field_mse, + standardized_field, + ) + + torch.manual_seed(6) + visual = torch.nn.functional.normalize(torch.randn(4, 8, 5), dim=-1) + noise = torch.nn.functional.normalize( + visual + 0.1 * torch.randn(4, 8, 5), dim=-1 + ) + visual_fields = standardized_field(visual) + text_fields = standardized_field(noise) + starts = torch.stack([torch.randperm(8) for _ in range(4)]) + final_cost, permutations = batched_descent( + text_fields.clone(), visual_fields, starts.clone(), sweeps=100 + ) + index = torch.arange(4) + permuted = text_fields[ + index[:, None, None], permutations[:, :, None], permutations[:, None, :] + ] + assert torch.allclose(field_mse(permuted, visual_fields), final_cost, atol=1e-10) + delta = batched_transposition_delta(permuted, visual_fields) + assert (delta >= -1e-9).all() + + +def test_ricci_flow_returns_finite_symmetric_geometry(): + import numpy as np + + from worldalign.ricci_control import ollivier_ricci_apsp + + rng = np.random.default_rng(4) + points = rng.normal(size=(24, 3)) + distance = np.linalg.norm(points[:, None] - points[None, :], axis=-1) + + class FlowArgs: + neighbors = 5 + lazy_alpha = 0.5 + flow_step = 0.4 + + flowed = ollivier_ricci_apsp(distance, FlowArgs(), iterations=3) + assert np.isfinite(flowed).all() + assert np.allclose(flowed, flowed.T, atol=1e-9) + assert (np.diag(flowed) == 0).all() + + +def test_paired_relation_field_has_less_energy_than_permutation(): + generator = torch.Generator().manual_seed(7) + visual = torch.randn(32, 12, generator=generator) + paired = visual @ torch.randn(12, 20, generator=generator) + visual_relation, visual_standardized = standardized_relation(visual) + paired_energy = relation_field_energy( + visual_relation, visual_standardized, paired + )[0] + permutation = torch.randperm(32, generator=generator) + shuffled_energy = relation_field_energy( + visual_relation, visual_standardized, paired[permutation] + )[0] + assert paired_energy < shuffled_energy + + +def test_tier0_factor_vector_is_one_hot_per_factor(): + from worldalign.tier0_pipeline import factor_vector + + vector = factor_vector(3, 2, 0, classes=10) + assert vector.shape == (17,) + assert vector[:10].sum() == 1.0 and vector[3] == 1.0 + assert vector[10:14].sum() == 1.0 and vector[11] == 1.0 + assert vector[14:].sum() == 1.0 and vector[14] == 1.0 + + +def test_watershed_splits_touching_objects(): + from worldalign.tier0_pipeline import extract_objects, group_objects + + image = torch.full((3, 64, 64), 0.1) + # Two same-coloured discs whose supports touch at a single column. + ys, xs = torch.meshgrid(torch.arange(64), torch.arange(64), indexing="ij") + for centre in (22, 38): + disc = ((ys - 32) ** 2 + (xs - centre) ** 2) <= 8**2 + image[:, disc] = torch.tensor([0.9, 0.2, 0.2])[:, None] + objects = extract_objects(image, peak_distance=4) + assert len(objects) == 2 + groups = group_objects(objects, threshold=0.35) + assert len(groups) == 1 and groups[0]["members"] == 2 + + +def test_caption_encoding_matches_declared_factors(): + from worldalign.tier0_pipeline import encode_caption + + ranks = {"red": 0, "blue": 1} + states = encode_caption("there are three small red circles, a blue square.", ranks, 2) + assert states.shape == (2, 9) + assert states[0][0] > 0 and states[0][2 + 2] > 0 and states[0][2 + 4 + 0] > 0 + assert states[1][1] > 0 and states[1][2 + 0] > 0 and states[1][2 + 4 + 1] > 0 |