tests/test_context_subset_consistency.py

import pytest
import torch

from mcss.context_subset_consistency import (
    compute_context_subset_geometry,
    drop_context_view,
    ramped_context_subset_geometry_weight,
)
from mcss.geometry import look_at, make_intrinsics
from mcss.measurements import FixedMeasurementRenderer
from mcss.types import Cameras, SceneState, StateEvidence


def _four_context_cameras() -> Cameras:
    intrinsics = make_intrinsics((4, 4), 35.0).view(1, 1, 3, 3).repeat(1, 4, 1, 1)
    intrinsics[0, :, 0, 0] += torch.arange(4)
    c2w = torch.eye(4).view(1, 1, 4, 4).repeat(1, 4, 1, 1)
    c2w[0, :, 0, 3] = torch.arange(4)
    return Cameras(intrinsics, c2w, (4, 4))


def _target_cameras() -> Cameras:
    intrinsics = make_intrinsics((4, 4), 35.0).view(1, 1, 3, 3)
    c2w = look_at(torch.tensor([0.0, 0.0, -3.0]), torch.zeros(3)).view(1, 1, 4, 4)
    return Cameras(intrinsics, c2w, (4, 4))


def _evidence_state(
    *,
    residual: torch.Tensor,
    confidence: torch.Tensor,
    provenance: torch.Tensor,
) -> SceneState:
    unknown = 1.0 - confidence
    gate = 0.1 + 0.9 * unknown
    base_density = 4.0 * confidence - 2.0
    density = base_density + gate * residual
    color = torch.full((1, 3, *confidence.shape[-3:]), 0.5)
    evidence = StateEvidence(
        confidence=confidence,
        unknown_probability=unknown,
        completion_gate=gate,
        provenance=provenance,
        base_density_logits=base_density,
        base_color=color,
        density_residual=residual,
        color_logit_residual=torch.zeros_like(color),
    )
    return SceneState(
        density_logits=density,
        color=color,
        log_variance=torch.full_like(density, -3.0),
        bounds=torch.tensor([[[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]]),
        evidence=evidence,
    )


def test_drop_context_view_preserves_retained_rgb_camera_order() -> None:
    rgb = torch.arange(4, dtype=torch.float32).view(1, 4, 1, 1, 1).expand(1, 4, 3, 4, 4)
    cameras = _four_context_cameras()

    subset = drop_context_view(rgb, cameras, dropped_index=2)

    assert subset.dropped_index == 2
    assert subset.retained_indices == (0, 1, 3)
    torch.testing.assert_close(subset.rgb[:, :, 0, 0, 0], torch.tensor([[0.0, 1.0, 3.0]]))
    torch.testing.assert_close(
        subset.cameras.intrinsics,
        cameras.intrinsics[:, [0, 1, 3]],
    )
    torch.testing.assert_close(subset.cameras.c2w, cameras.c2w[:, [0, 1, 3]])


@pytest.mark.parametrize("dropped_index", [-1, 4])
def test_drop_context_view_rejects_invalid_index(dropped_index: int) -> None:
    rgb = torch.zeros(1, 4, 3, 4, 4)

    with pytest.raises(ValueError, match="dropped_index"):
        drop_context_view(rgb, _four_context_cameras(), dropped_index=dropped_index)


def test_context_subset_weight_uses_optimizer_update_warmup() -> None:
    assert ramped_context_subset_geometry_weight(0.0, optimizer_step=1, warmup_steps=4) == 0.0
    assert ramped_context_subset_geometry_weight(0.1, optimizer_step=1, warmup_steps=4) == 0.025
    assert ramped_context_subset_geometry_weight(0.1, optimizer_step=4, warmup_steps=4) == 0.1
    assert ramped_context_subset_geometry_weight(0.1, optimizer_step=8, warmup_steps=4) == 0.1
    assert ramped_context_subset_geometry_weight(0.1, optimizer_step=1, warmup_steps=0) == 0.1


def test_context_subset_geometry_updates_only_subset_residual() -> None:
    shape = (1, 1, 4, 4, 4)
    full_confidence = torch.full(shape, 0.5, requires_grad=True)
    subset_confidence = torch.full(shape, 0.45, requires_grad=True)
    full_residual = torch.zeros(shape, requires_grad=True)
    subset_residual = torch.full(shape, 1.5, requires_grad=True)
    full_provenance = torch.full((1, 4, 4, 4, 4), 0.25)
    subset_provenance = torch.full((1, 3, 4, 4, 4), 1.0 / 3.0)
    full_state = _evidence_state(
        residual=full_residual,
        confidence=full_confidence,
        provenance=full_provenance,
    )
    subset_state = _evidence_state(
        residual=subset_residual,
        confidence=subset_confidence,
        provenance=subset_provenance,
    )
    renderer = FixedMeasurementRenderer(n_samples=12, ray_chunk_size=32)
    target_cameras = _target_cameras()
    full_predictions = renderer(full_state, target_cameras, {"depth", "visibility"})

    result = compute_context_subset_geometry(
        full_state,
        subset_state,
        full_predictions,
        target_cameras,
        renderer,
        residual_saturation_threshold=3.8,
    )

    assert torch.isfinite(result.loss)
    assert result.loss > 0.0
    assert set(result.terms) == {"depth", "visibility"}
    assert result.diagnostics["shared_coverage"] == 1.0
    assert 0.0 < result.diagnostics["overlap_weight_mean"] < 1.0
    assert "full_residual_abs_p50" in result.audit_diagnostics
    assert len(result.per_sample_diagnostics) == 1
    result.loss.backward()
    assert subset_residual.grad is not None
    assert subset_residual.grad.abs().sum() > 0.0
    assert full_residual.grad is None
    assert full_confidence.grad is None
    assert subset_confidence.grad is None


def test_context_subset_geometry_zero_overlap_is_finite_zero() -> None:
    shape = (1, 1, 4, 4, 4)
    full_state = _evidence_state(
        residual=torch.zeros(shape),
        confidence=torch.full(shape, 0.5),
        provenance=torch.full((1, 4, 4, 4, 4), 0.25),
    )
    subset_residual = torch.full(shape, 3.0, requires_grad=True)
    subset_state = _evidence_state(
        residual=subset_residual,
        confidence=torch.zeros(shape),
        provenance=torch.zeros(1, 3, 4, 4, 4),
    )
    renderer = FixedMeasurementRenderer(n_samples=12, ray_chunk_size=32)
    target_cameras = _target_cameras()
    full_predictions = renderer(full_state, target_cameras, {"depth", "visibility"})

    result = compute_context_subset_geometry(
        full_state,
        subset_state,
        full_predictions,
        target_cameras,
        renderer,
        residual_saturation_threshold=3.8,
    )

    assert result.diagnostics["shared_coverage"] == 0.0
    assert result.diagnostics["overlap_weight_mean"] == 0.0
    torch.testing.assert_close(result.loss, torch.zeros_like(result.loss))
    result.loss.backward()
    assert subset_residual.grad is not None
    torch.testing.assert_close(subset_residual.grad, torch.zeros_like(subset_residual.grad))


def test_context_subset_geometry_skips_expensive_audit_when_not_logging() -> None:
    shape = (1, 1, 4, 4, 4)
    full_state = _evidence_state(
        residual=torch.zeros(shape),
        confidence=torch.full(shape, 0.5),
        provenance=torch.full((1, 4, 4, 4, 4), 0.25),
    )
    subset_state = _evidence_state(
        residual=torch.ones(shape, requires_grad=True),
        confidence=torch.full(shape, 0.45),
        provenance=torch.full((1, 3, 4, 4, 4), 1.0 / 3.0),
    )
    renderer = FixedMeasurementRenderer(n_samples=12, ray_chunk_size=32)
    target_cameras = _target_cameras()
    full_predictions = renderer(full_state, target_cameras, {"depth", "visibility"})

    result = compute_context_subset_geometry(
        full_state,
        subset_state,
        full_predictions,
        target_cameras,
        renderer,
        residual_saturation_threshold=3.8,
        collect_audit_diagnostics=False,
    )

    assert torch.isfinite(result.loss)
    assert result.audit_diagnostics == {}
    assert result.per_sample_diagnostics == ()