import pytest import torch from mcss.engine import _state_diagnostics, _state_regularization from mcss.geometry import look_at, make_intrinsics from mcss.model.evidence_encoder import ( _candidate_profile_telemetry, _cross_view_evidence, _normalized_center_peakness, _positive_candidate_excess, _surface_peakness, _surface_peakness_with_telemetry, _trilinear_splat_candidates, ) from mcss.model.system import ModelConfig, build_model from mcss.types import ( Cameras, DualEvidence, SceneState, StateAppearance, StateEvidence, TransportEvidence, ) def _model_inputs() -> tuple[torch.Tensor, Cameras, Cameras, torch.Tensor]: torch.manual_seed(13) context_rgb = torch.rand(1, 2, 3, 8, 8) intrinsics = make_intrinsics((8, 8), 60.0) context_poses = torch.stack( ( look_at(torch.tensor([0.0, 0.0, -3.0]), torch.zeros(3)), look_at(torch.tensor([2.0, 0.0, -2.0]), torch.zeros(3)), ) ) target_pose = look_at(torch.tensor([-2.0, 0.0, -2.0]), torch.zeros(3)) context_cameras = Cameras( intrinsics.expand(1, 2, 3, 3).clone(), context_poses.unsqueeze(0), (8, 8) ) target_cameras = Cameras( intrinsics.view(1, 1, 3, 3), target_pose.view(1, 1, 4, 4), (8, 8) ) bounds = torch.tensor([[[-1.2, -1.2, -1.2], [1.2, 1.2, 1.2]]]) return context_rgb, context_cameras, target_cameras, bounds def test_cross_view_evidence_prefers_agreement_and_masks_invalid_views() -> None: features = torch.tensor( [[[[1.0, 0.0], [1.0, 0.0]], [[1.0, 0.0], [-1.0, 0.0]]]] ) rgb = torch.tensor( [[[[0.2, 0.3, 0.4], [0.1, 0.1, 0.1]], [[0.2, 0.3, 0.4], [0.9, 0.9, 0.9]]]] ) valid = torch.ones(1, 2, 2, dtype=torch.bool) confidence, provenance = _cross_view_evidence(features, rgb, valid, temperature=0.1) assert confidence.shape == (1, 2, 1) assert provenance.shape == (1, 2, 2, 1) assert confidence[0, 0, 0] > confidence[0, 1, 0] torch.testing.assert_close(provenance.sum(dim=1), torch.ones(1, 2, 1)) one_view_valid = torch.tensor([[[True], [False]]]) confidence, provenance = _cross_view_evidence( features[:, :, :1], rgb[:, :, :1], one_view_valid, temperature=0.1 ) torch.testing.assert_close(confidence, torch.zeros_like(confidence)) assert provenance[0, 1, 0, 0] == 0.0 assert provenance[0, 0, 0, 0] == 1.0 def test_cross_view_evidence_cannot_optimize_its_own_confidence() -> None: features = torch.tensor( [[[[1.0, 0.0]], [[0.8, 0.2]]]], requires_grad=True, ) rgb = torch.tensor([[[[0.2, 0.3, 0.4]], [[0.2, 0.3, 0.4]]]]) valid = torch.ones(1, 2, 1, dtype=torch.bool) confidence, provenance = _cross_view_evidence(features, rgb, valid, temperature=0.1) assert not confidence.requires_grad assert not provenance.requires_grad def test_center_peakness_is_zero_for_flat_valid_profiles() -> None: scores = torch.zeros(1, 3, 5) valid = torch.ones_like(scores, dtype=torch.bool) peakness, contributes = _normalized_center_peakness( scores, valid, temperature=0.05, ) torch.testing.assert_close(peakness, torch.zeros_like(peakness), atol=1e-6, rtol=0) assert contributes.all() def test_center_peakness_uses_valid_candidate_count_and_rejects_invalid_centers() -> None: scores = torch.tensor( [[[0.0, 0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0]]] ) valid = torch.tensor( [[[True, False, True, True, False], [True, True, False, True, True]]] ) peakness, contributes = _normalized_center_peakness( scores, valid, temperature=0.05, ) assert peakness[0, 0] > 0.99 assert contributes.tolist() == [[True, False]] assert peakness[0, 1] == 0.0 def test_candidate_profile_telemetry_distinguishes_flat_and_shifted_winners() -> None: scores = torch.tensor( [ [0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.1, 0.2, 0.8, 0.0], ] ) valid = torch.ones_like(scores, dtype=torch.bool) observer_count = torch.full_like(scores, 3, dtype=torch.int64) report = _candidate_profile_telemetry( scores, valid, observer_count, offsets=torch.tensor([-2, -1, 0, 1, 2]), temperature=0.05, appearance_confidence=torch.tensor([0.0, 1.0]), observer_capacity=3, ) assert report["profiles"] == { "total": 2, "contributing": 2, "contributing_fraction": 1.0, } assert report["valid_candidate_count_histogram"] == {"5": 2} assert report["observer_count_histogram"] == {"3": 10} assert report["competition"]["center_argmax_inclusive_rate"] == pytest.approx(0.5) assert report["competition"]["center_unique_argmax_rate"] == pytest.approx(0.0) assert report["competition"]["center_argmax_fractional_rate"] == pytest.approx(0.1) assert report["competition"]["chance_argmax_rate"] == pytest.approx(0.2) assert report["competition"]["center_argmax_excess_over_chance"] == pytest.approx(-0.1) assert report["competition"]["unique_winner_fraction"] == pytest.approx(0.5) assert report["competition"]["winner_offset_histogram"] == {"1": 1} assert report["competition"]["noncenter_positive_excess_fraction"] == pytest.approx(1.0) assert report["competition"]["zero_positive_excess_profile_count"] == 1 assert report["competition"]["score_range"]["p50"] == pytest.approx(0.4) assert report["competition"]["softmax_entropy_normalized"]["max"] == pytest.approx(1.0) assert report["strata"]["appearance_positive"]["profiles"] == 1 assert report["strata"]["appearance_positive"]["competition"][ "winner_offset_histogram" ] == {"1": 1} assert report["strata"]["appearance_top_decile"]["profiles"] == 1 assert report["counterfactuals"]["valid_observer_mean"]["competition"][ "center_argmax_fractional_rate" ] == pytest.approx(0.1) assert report["counterfactuals"]["min_two_observers"]["profiles"][ "contributing" ] == 2 def test_candidate_profile_telemetry_handles_no_contributing_profiles() -> None: scores = torch.zeros(2, 5) valid = torch.zeros_like(scores, dtype=torch.bool) observer_count = torch.zeros_like(scores, dtype=torch.int64) report = _candidate_profile_telemetry( scores, valid, observer_count, offsets=torch.tensor([-2, -1, 0, 1, 2]), temperature=0.05, ) assert report["profiles"]["contributing"] == 0 assert report["competition"] == {} def test_surface_peakness_telemetry_preserves_the_production_result() -> None: context_rgb, context_cameras, _, bounds = _model_inputs() appearance_confidence = torch.linspace(0.0, 1.0, 4**3).reshape(1, 1, 4, 4, 4) kwargs = { "evidence_temperature": 0.1, "surface_peak_temperature": 0.05, "chunk_size": 16, } expected = _surface_peakness( context_rgb, context_cameras, bounds, (4, 4, 4), **kwargs, ) actual, report = _surface_peakness_with_telemetry( context_rgb, context_cameras, bounds, (4, 4, 4), appearance_confidence=appearance_confidence, **kwargs, ) torch.testing.assert_close(actual, expected) assert report["profiles"]["total"] == 2 * 4 * 4 * 4 assert sum(report["valid_candidate_count_histogram"].values()) == 2 * 4 * 4 * 4 assert report["surface_peakness"]["count"] == 4 * 4 * 4 assert "appearance_positive" in report["strata"] assert "appearance_top_decile" in report["strata"] def test_trilinear_splat_places_exact_center_mass_in_one_voxel() -> None: points = torch.tensor([[[[0.5, 0.5, 0.5]]]]) values = torch.tensor([[[2.0]]]) valid = torch.ones_like(values, dtype=torch.bool) bounds = torch.tensor([[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]]) volume = _trilinear_splat_candidates(points, values, valid, bounds, (2, 2, 2)) expected = torch.zeros(1, 8) expected[0, 0] = 2.0 torch.testing.assert_close(volume, expected) def test_trilinear_splat_splits_midpoint_mass_across_eight_voxels() -> None: points = torch.tensor([[[[1.0, 1.0, 1.0]]]]) values = torch.ones(1, 1, 1) valid = torch.ones_like(values, dtype=torch.bool) bounds = torch.tensor([[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]]) volume = _trilinear_splat_candidates(points, values, valid, bounds, (2, 2, 2)) torch.testing.assert_close(volume, torch.full((1, 8), 0.125)) assert volume.sum().item() == pytest.approx(1.0) def test_trilinear_splat_renormalizes_mass_at_the_aabb_boundary() -> None: points = torch.tensor([[[[0.0, 0.0, 0.0]]]]) values = torch.tensor([[[3.0]]]) valid = torch.ones_like(values, dtype=torch.bool) bounds = torch.tensor([[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]]) volume = _trilinear_splat_candidates(points, values, valid, bounds, (2, 2, 2)) expected = torch.zeros(1, 8) expected[0, 0] = 3.0 torch.testing.assert_close(volume, expected) assert volume.sum().item() == pytest.approx(3.0) def test_trilinear_splat_reuses_weights_for_multiple_value_channels() -> None: points = torch.tensor([[[[0.5, 0.5, 0.5]]]]) values = torch.tensor([[[[2.0, 3.0]]]]) valid = torch.ones(1, 1, 1, dtype=torch.bool) bounds = torch.tensor([[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]]) volume = _trilinear_splat_candidates(points, values, valid, bounds, (2, 2, 2)) expected = torch.zeros(1, 2, 8) expected[0, :, 0] = torch.tensor([2.0, 3.0]) torch.testing.assert_close(volume, expected) def test_positive_candidate_excess_is_zero_for_flat_and_tracks_offcenter_peak() -> None: scores = torch.tensor( [ [0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0], ] ) valid = torch.ones_like(scores, dtype=torch.bool) excess, contributes = _positive_candidate_excess(scores, valid, temperature=0.05) torch.testing.assert_close(excess[0], torch.zeros(5), atol=1e-6, rtol=0) assert excess[1, 3] > 0.99 assert excess[1, 2] == 0.0 assert contributes.tolist() == [True, True] def test_positive_excess_is_splatted_to_the_offcenter_candidate_location() -> None: points = torch.tensor( [ [ [ [0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [1.5, 0.5, 0.5], [0.5, 0.5, 0.5], ] ] ] ) scores = torch.tensor([[[0.0, 0.0, 0.0, 1.0, 0.0]]]) valid = torch.ones_like(scores, dtype=torch.bool) bounds = torch.tensor([[[0.0, 0.0, 0.0], [3.0, 1.0, 1.0]]]) excess, contributes = _positive_candidate_excess( scores, valid, temperature=0.05 ) transported = _trilinear_splat_candidates( points, excess, valid & contributes[..., None], bounds, (1, 1, 3), ) assert transported[0, 0].item() == pytest.approx(0.0, abs=1e-6) assert transported[0, 1].item() > 0.99 def _transport_model() -> torch.nn.Module: return build_model( ModelConfig( mode="fixed", state_architecture="dual_evidence_transport", voxel_resolution=4, image_feature_dim=8, state_feature_dim=8, refinement_blocks=1, evidence_temperature=0.1, surface_peak_temperature=0.05, observed_residual_floor=0.2, completion_residual_scale=1.5, appearance_resolution_scale=2, n_samples=4, ray_chunk_size=64, ) ) def test_transport_architecture_keeps_parameters_and_requires_two_other_observers() -> None: v5 = build_model( ModelConfig( mode="fixed", state_architecture="evidence_residual", voxel_resolution=4, image_feature_dim=8, state_feature_dim=8, refinement_blocks=1, appearance_resolution_scale=2, n_samples=4, ray_chunk_size=64, ) ) model = _transport_model() output = model(*_model_inputs()) evidence = output.state.transport_evidence assert set(model.state_dict()) == set(v5.state_dict()) assert evidence is not None assert output.state.dual_evidence is None assert output.state.evidence is None torch.testing.assert_close( evidence.surface_localization_support, evidence.transported_appearance_confidence * evidence.surface_peakness, ) torch.testing.assert_close( evidence.surface_localization_support, torch.zeros_like(evidence.surface_localization_support), ) torch.testing.assert_close(output.state.density_logits, evidence.density_logits) torch.testing.assert_close(output.state.color, evidence.color) assert not evidence.surface_localization_support.requires_grad def test_transport_evidence_is_invariant_to_context_view_permutation() -> None: model = _transport_model().eval() context_rgb, context_cameras, target_cameras, bounds = _model_inputs() context_rgb = torch.cat((context_rgb, context_rgb[:, :1]), dim=1) context_cameras = Cameras( torch.cat((context_cameras.intrinsics, context_cameras.intrinsics[:, :1]), dim=1), torch.cat((context_cameras.c2w, context_cameras.c2w[:, :1]), dim=1), context_cameras.image_size, ) permutation = torch.tensor([2, 0, 1]) original = model(context_rgb, context_cameras, target_cameras, bounds).state permuted = model( context_rgb[:, permutation], context_cameras.select_views(permutation), target_cameras, bounds, ).state assert original.transport_evidence is not None assert permuted.transport_evidence is not None torch.testing.assert_close( original.transport_evidence.surface_localization_support, permuted.transport_evidence.surface_localization_support, atol=1e-5, rtol=1e-5, ) torch.testing.assert_close( original.transport_evidence.transport_coverage, permuted.transport_evidence.transport_coverage, atol=1e-5, rtol=1e-5, ) torch.testing.assert_close( original.transport_evidence.provenance[:, permutation], permuted.transport_evidence.provenance, atol=1e-5, rtol=1e-5, ) def test_transport_evidence_cannot_change_across_an_optimizer_step() -> None: model = _transport_model() context_rgb, context_cameras, target_cameras, bounds = _model_inputs() context_rgb = torch.cat((context_rgb, context_rgb[:, :1]), dim=1) context_cameras = Cameras( torch.cat((context_cameras.intrinsics, context_cameras.intrinsics[:, :1]), dim=1), torch.cat((context_cameras.c2w, context_cameras.c2w[:, :1]), dim=1), context_cameras.image_size, ) inputs = (context_rgb, context_cameras, target_cameras, bounds) optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) before = model(*inputs) assert before.state.transport_evidence is not None frozen = { "localization": before.state.transport_evidence.surface_localization_support.clone(), "transported_appearance": ( before.state.transport_evidence.transported_appearance_confidence.clone() ), "peakness": before.state.transport_evidence.surface_peakness.clone(), "coverage": before.state.transport_evidence.transport_coverage.clone(), "appearance": before.state.transport_evidence.appearance_confidence.clone(), "provenance": before.state.transport_evidence.provenance.clone(), } loss = sum(value.mean() for value in before.predictions.values()) loss.backward() optimizer.step() after = model(*inputs) assert after.state.transport_evidence is not None torch.testing.assert_close( after.state.transport_evidence.surface_localization_support, frozen["localization"], ) torch.testing.assert_close( after.state.transport_evidence.transported_appearance_confidence, frozen["transported_appearance"], ) torch.testing.assert_close( after.state.transport_evidence.surface_peakness, frozen["peakness"], ) torch.testing.assert_close( after.state.transport_evidence.transport_coverage, frozen["coverage"], ) torch.testing.assert_close( after.state.transport_evidence.appearance_confidence, frozen["appearance"], ) torch.testing.assert_close( after.state.transport_evidence.provenance, frozen["provenance"], ) def test_dual_evidence_architecture_keeps_v5_parameter_contract() -> None: common = dict( mode="fixed", voxel_resolution=4, image_feature_dim=8, state_feature_dim=8, refinement_blocks=1, appearance_resolution_scale=2, n_samples=4, ray_chunk_size=64, ) v5 = build_model(ModelConfig(state_architecture="evidence_residual", **common)) v6a = build_model( ModelConfig( state_architecture="dual_evidence", surface_peak_temperature=0.05, **common, ) ) assert set(v6a.state_dict()) == set(v5.state_dict()) load_result = v5.load_state_dict(v5.state_dict(), strict=True) assert load_result.missing_keys == [] assert load_result.unexpected_keys == [] def _dual_evidence_model() -> torch.nn.Module: return build_model( ModelConfig( mode="fixed", state_architecture="dual_evidence", voxel_resolution=4, image_feature_dim=8, state_feature_dim=8, refinement_blocks=1, evidence_temperature=0.1, surface_peak_temperature=0.05, observed_residual_floor=0.2, completion_residual_scale=1.5, appearance_resolution_scale=2, n_samples=4, ray_chunk_size=64, ) ) def test_dual_evidence_forward_uses_separate_density_and_appearance_gates() -> None: model = _dual_evidence_model() output = model(*_model_inputs()) evidence = output.state.dual_evidence assert evidence is not None assert output.state.evidence is None assert output.state.appearance is not None assert evidence.surface_localization_support.shape == (1, 1, 4, 4, 4) assert evidence.appearance_confidence.shape == (1, 1, 4, 4, 4) assert evidence.surface_peakness.shape == (1, 1, 4, 4, 4) assert not evidence.surface_localization_support.requires_grad assert not evidence.appearance_confidence.requires_grad assert not evidence.surface_peakness.requires_grad torch.testing.assert_close( evidence.surface_localization_support, evidence.appearance_confidence * evidence.surface_peakness, ) torch.testing.assert_close(output.state.density_logits, evidence.density_logits) torch.testing.assert_close(output.state.color, evidence.color) assert (evidence.localization_completion_gate >= 0.2).all() assert (evidence.appearance_completion_gate >= 0.2).all() assert all(torch.isfinite(value).all() for value in output.predictions.values()) loss = sum(value.mean() for value in output.predictions.values()) loss.backward() gradients = [parameter.grad for parameter in model.parameters() if parameter.requires_grad] assert any(gradient is not None for gradient in gradients) assert all(torch.isfinite(gradient).all() for gradient in gradients if gradient is not None) def test_dual_evidence_is_invariant_to_context_view_permutation() -> None: model = _dual_evidence_model().eval() context_rgb, context_cameras, target_cameras, bounds = _model_inputs() permutation = torch.tensor([1, 0]) original = model(context_rgb, context_cameras, target_cameras, bounds).state permuted = model( context_rgb[:, permutation], context_cameras.select_views(permutation), target_cameras, bounds, ).state assert original.dual_evidence is not None assert permuted.dual_evidence is not None torch.testing.assert_close( original.dual_evidence.surface_localization_support, permuted.dual_evidence.surface_localization_support, atol=1e-5, rtol=1e-5, ) torch.testing.assert_close( original.dual_evidence.appearance_confidence, permuted.dual_evidence.appearance_confidence, atol=1e-5, rtol=1e-5, ) torch.testing.assert_close( original.dual_evidence.surface_peakness, permuted.dual_evidence.surface_peakness, atol=1e-5, rtol=1e-5, ) torch.testing.assert_close( original.dual_evidence.provenance[:, permutation], permuted.dual_evidence.provenance, atol=1e-5, rtol=1e-5, ) def test_dual_evidence_cannot_change_across_an_optimizer_step() -> None: model = _dual_evidence_model() inputs = _model_inputs() optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) before = model(*inputs) assert before.state.dual_evidence is not None frozen = { "localization": before.state.dual_evidence.surface_localization_support.clone(), "appearance": before.state.dual_evidence.appearance_confidence.clone(), "peakness": before.state.dual_evidence.surface_peakness.clone(), "provenance": before.state.dual_evidence.provenance.clone(), } loss = sum(value.mean() for value in before.predictions.values()) loss.backward() optimizer.step() after = model(*inputs) assert after.state.dual_evidence is not None torch.testing.assert_close( after.state.dual_evidence.surface_localization_support, frozen["localization"], ) torch.testing.assert_close( after.state.dual_evidence.appearance_confidence, frozen["appearance"], ) torch.testing.assert_close(after.state.dual_evidence.surface_peakness, frozen["peakness"]) torch.testing.assert_close(after.state.dual_evidence.provenance, frozen["provenance"]) def test_evidence_residual_state_is_auditable_and_reconstructable() -> None: config = ModelConfig( mode="fixed", state_architecture="evidence_residual", voxel_resolution=8, image_feature_dim=8, state_feature_dim=8, refinement_blocks=1, evidence_temperature=0.1, observed_residual_floor=0.2, completion_residual_scale=1.5, n_samples=8, ray_chunk_size=128, ) model = build_model(config) output = model(*_model_inputs()) evidence = output.state.evidence assert evidence is not None assert output.state.appearance is None assert evidence.confidence.shape == (1, 1, 8, 8, 8) assert evidence.provenance.shape == (1, 2, 8, 8, 8) assert (evidence.completion_gate >= 0.2).all() assert (evidence.completion_gate <= 1.0).all() assert evidence.density_residual.abs().max() <= 1.5 assert evidence.color_logit_residual.abs().max() <= 1.5 assert evidence.base_density_logits.min() >= -2.0001 assert evidence.base_density_logits.max() <= 2.0001 empty_provenance = evidence.provenance.sum(dim=1, keepdim=True) == 0 if empty_provenance.any(): empty_color = empty_provenance.expand_as(evidence.base_color) torch.testing.assert_close( evidence.base_color[empty_color], torch.full_like(evidence.base_color[empty_color], 0.5), ) reconstructed_density = ( evidence.base_density_logits + evidence.completion_gate * evidence.density_residual ) base_logits = torch.logit(evidence.base_color.clamp(1e-4, 1.0 - 1e-4)) reconstructed_color = torch.sigmoid( base_logits + evidence.completion_gate * evidence.color_logit_residual ) torch.testing.assert_close(output.state.density_logits, reconstructed_density) torch.testing.assert_close(output.state.color, reconstructed_color) assert all(torch.isfinite(value).all() for value in output.predictions.values()) loss = sum(value.mean() for value in output.predictions.values()) loss.backward() gradients = [parameter.grad for parameter in model.parameters() if parameter.requires_grad] assert any(gradient is not None for gradient in gradients) assert all(torch.isfinite(gradient).all() for gradient in gradients if gradient is not None) def test_highres_appearance_is_auditable_reconstructable_and_differentiable() -> None: config = ModelConfig( mode="fixed", state_architecture="evidence_residual", voxel_resolution=8, image_feature_dim=8, state_feature_dim=8, refinement_blocks=1, evidence_temperature=0.1, observed_residual_floor=0.2, completion_residual_scale=1.5, appearance_resolution_scale=2, n_samples=8, ray_chunk_size=128, ) model = build_model(config) output = model(*_model_inputs()) appearance = output.state.appearance assert appearance is not None assert output.state.spatial_shape == (8, 8, 8) assert output.state.appearance_shape == (16, 16, 16) assert appearance.confidence.shape == (1, 1, 16, 16, 16) assert appearance.provenance.shape == (1, 2, 16, 16, 16) assert appearance.base_color.shape == (1, 3, 16, 16, 16) assert (appearance.completion_gate >= 0.2).all() assert (appearance.completion_gate <= 1.0).all() assert appearance.color_logit_residual.abs().max() <= 1.5 torch.testing.assert_close( appearance.color_logit_residual, torch.zeros_like(appearance.color_logit_residual), ) expected_color = torch.sigmoid(torch.logit(appearance.base_color.clamp(1e-4, 1 - 1e-4))) torch.testing.assert_close(appearance.color, expected_color) assert all(torch.isfinite(value).all() for value in output.predictions.values()) loss = sum(value.mean() for value in output.predictions.values()) loss.backward() gradients = [parameter.grad for parameter in model.parameters() if parameter.requires_grad] assert any(gradient is not None for gradient in gradients) assert all(torch.isfinite(gradient).all() for gradient in gradients if gradient is not None) def test_evidence_fields_are_predictor_independent_across_an_optimizer_step() -> None: config = ModelConfig( mode="fixed", state_architecture="evidence_residual", voxel_resolution=8, image_feature_dim=8, state_feature_dim=8, refinement_blocks=1, appearance_resolution_scale=2, n_samples=8, ray_chunk_size=128, ) model = build_model(config) inputs = _model_inputs() optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) before = model(*inputs) assert before.state.evidence is not None confidence = before.state.evidence.confidence.detach().clone() provenance = before.state.evidence.provenance.detach().clone() assert before.state.appearance is not None appearance_confidence = before.state.appearance.confidence.detach().clone() appearance_provenance = before.state.appearance.provenance.detach().clone() appearance_base_color = before.state.appearance.base_color.detach().clone() loss = sum(value.mean() for value in before.predictions.values()) loss.backward() optimizer.step() after = model(*inputs) assert after.state.evidence is not None assert after.state.appearance is not None torch.testing.assert_close(after.state.evidence.confidence, confidence) torch.testing.assert_close(after.state.evidence.provenance, provenance) torch.testing.assert_close(after.state.appearance.confidence, appearance_confidence) torch.testing.assert_close(after.state.appearance.provenance, appearance_provenance) torch.testing.assert_close(after.state.appearance.base_color, appearance_base_color) def test_evidence_residual_regularization_preserves_legacy_and_penalizes_rewriting() -> None: shape = (1, 1, 4, 4, 4) bounds = torch.tensor([[[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]]) density = torch.zeros(shape) color = torch.full((1, 3, 4, 4, 4), 0.5) variance = torch.zeros(shape) legacy = SceneState(density, color, variance, bounds) legacy_unweighted = _state_regularization(legacy, evidence_residual_weight=0.0) legacy_weighted = _state_regularization(legacy, evidence_residual_weight=1.0) torch.testing.assert_close(legacy_unweighted, legacy_weighted) assert _state_diagnostics(legacy) == {} confidence = torch.full(shape, 0.8) unknown = 1.0 - confidence evidence = StateEvidence( confidence=confidence, unknown_probability=unknown, completion_gate=unknown, provenance=torch.full((1, 2, 4, 4, 4), 0.5), base_density_logits=density, base_color=color, density_residual=torch.full(shape, 0.25), color_logit_residual=torch.full_like(color, 0.5), ) state = SceneState(density, color, variance, bounds, evidence=evidence) unweighted = _state_regularization(state, evidence_residual_weight=0.0) weighted = _state_regularization(state, evidence_residual_weight=0.01) diagnostics = _state_diagnostics(state) assert weighted > unweighted assert diagnostics["state/evidence_mean"] == pytest.approx(0.8) assert diagnostics["state/unknown_mean"] == pytest.approx(0.2) assert diagnostics["state/completion_gate_mean"] == pytest.approx(0.2) assert diagnostics["state/density_residual_abs"] == pytest.approx(0.25) assert diagnostics["state/color_residual_abs"] == pytest.approx(0.5) appearance_confidence = torch.full((1, 1, 8, 8, 8), 0.75) appearance = StateAppearance( confidence=appearance_confidence, unknown_probability=1.0 - appearance_confidence, completion_gate=torch.full_like(appearance_confidence, 0.325), provenance=torch.full((1, 2, 8, 8, 8), 0.5), base_color=torch.full((1, 3, 8, 8, 8), 0.5), color_logit_residual=torch.ones(1, 3, 8, 8, 8), ) appearance_state = SceneState( density, color, variance, bounds, evidence=evidence, appearance=appearance, ) appearance_weighted = _state_regularization( appearance_state, evidence_residual_weight=0.01 ) appearance_diagnostics = _state_diagnostics(appearance_state) expected_rewrite = torch.tensor(0.01 * (0.8 * 0.25 + 0.75 * 1.0)) torch.testing.assert_close(appearance_weighted, unweighted + expected_rewrite) assert appearance_diagnostics["state/appearance_evidence_mean"] == pytest.approx(0.75) assert appearance_diagnostics["state/appearance_unknown_mean"] == pytest.approx(0.25) assert appearance_diagnostics["state/appearance_completion_gate_mean"] == pytest.approx(0.325) assert appearance_diagnostics["state/appearance_color_residual_abs"] == pytest.approx(1.0) assert appearance_diagnostics["state/appearance_resolution_scale"] == pytest.approx(2.0) def test_dual_evidence_regularization_uses_localization_for_density_and_appearance_for_color( ) -> None: shape = (1, 1, 4, 4, 4) localization = torch.full(shape, 0.2) appearance_confidence = torch.full(shape, 0.8) peakness = localization / appearance_confidence base_color = torch.full((1, 3, 4, 4, 4), 0.5) dual = DualEvidence( surface_localization_support=localization, appearance_confidence=appearance_confidence, surface_peakness=peakness, localization_unknown_probability=1.0 - localization, appearance_unknown_probability=1.0 - appearance_confidence, localization_completion_gate=1.0 - localization, appearance_completion_gate=1.0 - appearance_confidence, provenance=torch.full((1, 2, 4, 4, 4), 0.5), base_density_logits=torch.zeros(shape), base_color=base_color, density_residual=torch.full(shape, 0.5), color_logit_residual=torch.ones_like(base_color), ) state = SceneState( density_logits=dual.density_logits, color=dual.color, log_variance=torch.zeros(shape), bounds=torch.tensor([[[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]]), dual_evidence=dual, ) unweighted = _state_regularization(state, evidence_residual_weight=0.0) weighted = _state_regularization(state, evidence_residual_weight=0.01) diagnostics = _state_diagnostics(state) expected_rewrite = torch.tensor(0.01 * (0.2 * 0.5 + 0.8 * 1.0)) torch.testing.assert_close(weighted, unweighted + expected_rewrite) assert diagnostics["state/surface_localization_support_mean"] == pytest.approx(0.2) assert diagnostics["state/appearance_confidence_mean"] == pytest.approx(0.8) assert diagnostics["state/surface_peakness_mean"] == pytest.approx(0.25) assert diagnostics["state/surface_localization_support_p50"] == pytest.approx(0.2) assert diagnostics["state/appearance_confidence_p50"] == pytest.approx(0.8) assert diagnostics["state/surface_peakness_p50"] == pytest.approx(0.25) assert diagnostics["state/surface_peakness_nonzero_fraction"] == pytest.approx(1.0) assert diagnostics["state/localization_completion_gate_mean"] == pytest.approx(0.8) assert diagnostics["state/appearance_completion_gate_mean"] == pytest.approx(0.2) def test_transport_regularization_and_diagnostics_use_transported_localization() -> None: shape = (1, 1, 4, 4, 4) localization = torch.full(shape, 0.2) transported_appearance = torch.full(shape, 0.4) appearance_confidence = torch.full(shape, 0.8) base_color = torch.full((1, 3, 4, 4, 4), 0.5) transport = TransportEvidence( surface_localization_support=localization, transported_appearance_confidence=transported_appearance, surface_peakness=localization / transported_appearance, transport_coverage=torch.full(shape, 3.0), appearance_confidence=appearance_confidence, localization_unknown_probability=1.0 - localization, appearance_unknown_probability=1.0 - appearance_confidence, localization_completion_gate=1.0 - localization, appearance_completion_gate=1.0 - appearance_confidence, provenance=torch.full((1, 2, 4, 4, 4), 0.5), base_density_logits=torch.zeros(shape), base_color=base_color, density_residual=torch.full(shape, 0.5), color_logit_residual=torch.ones_like(base_color), ) state = SceneState( density_logits=transport.density_logits, color=transport.color, log_variance=torch.zeros(shape), bounds=torch.tensor([[[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]]), transport_evidence=transport, ) unweighted = _state_regularization(state, evidence_residual_weight=0.0) weighted = _state_regularization(state, evidence_residual_weight=0.01) diagnostics = _state_diagnostics(state) expected_rewrite = torch.tensor(0.01 * (0.2 * 0.5 + 0.8 * 1.0)) torch.testing.assert_close(weighted, unweighted + expected_rewrite) assert diagnostics["state/surface_localization_support_mean"] == pytest.approx(0.2) assert diagnostics["state/transported_appearance_confidence_mean"] == pytest.approx(0.4) assert diagnostics["state/surface_peakness_mean"] == pytest.approx(0.5) assert diagnostics["state/transport_coverage_mean"] == pytest.approx(3.0) assert diagnostics["state/transport_coverage_nonzero_fraction"] == pytest.approx(1.0) assert diagnostics["state/localization_density_rewrite_mean"] == pytest.approx(0.1)