tests/test_direct_capacity_optimization.py
"""Independent fixed-kernel and context-only direct optimization contracts."""
import ast
import inspect
import json
from dataclasses import replace
import pytest
import torch
from mcss.dynamic.types import clone_scene_state, hash_scene_state
from mcss.measurements import FixedMeasurementRenderer
from mcss.mechanism_pilot import direct_capacity_optimization as optimization
from mcss.mechanism_pilot.direct_capacity_contracts import ObservationBatch
from mcss.types import Cameras, SceneState
def fixture_batch(depth=True):
poses = torch.eye(4).repeat(1, 2, 1, 1)
poses[0, 1, 0, 3] = 0.2
k = torch.tensor([[4.0, 0.0, 2.0], [0.0, 4.0, 2.0], [0.0, 0.0, 1.0]])
cameras = Cameras(k.repeat(1, 2, 1, 1), poses, (5, 5))
bounds = torch.tensor([[-2.0, -2.0, 0.1], [2.0, 2.0, 4.0]])
target = optimization.DirectState(bounds, 4)
with torch.no_grad():
target.density.fill_(-0.4)
target.color_logits[:, 0].fill_(1.0)
target.color_logits[:, 1].fill_(-0.6)
rendered = FixedMeasurementRenderer(n_samples=16)(target.state(), cameras)
batch = ObservationBatch(
"synthetic",
(0, 3),
rendered["rgb"].detach(),
cameras,
rendered["depth"].detach() if depth else None,
"CONTEXT_RGBD" if depth else "CONTEXT_RGB_ONLY",
)
return batch, bounds
def flatten_image(image):
return image.permute(0, 1, 3, 4, 2).reshape(1, -1, image.shape[2])
@pytest.mark.parametrize("samples", [64, 128])
def test_cached_selected_rays_match_full_predictions_and_parameter_gradients(samples):
batch, bounds = fixture_batch()
model = optimization.DirectState(bounds, 4)
with torch.no_grad():
model.density.add_(torch.linspace(-0.1, 0.2, 64).reshape_as(model.density))
model.color_logits.add_(torch.linspace(-1.0, 1.0, 192).reshape_as(model.color_logits))
renderer = FixedMeasurementRenderer(n_samples=samples)
cache = optimization.cache_rays(batch, bounds)
# Both frames, boundary pixels and duplicate rays exercise flattening and accumulation.
selected = torch.tensor([0, 4, 12, 24, 25, 31, 49, 12])
full = renderer(model.state(), batch.cameras)
sampled = optimization.render_rays(renderer, model.state(), cache, selected)
for field in ("rgb", "depth", "visibility"):
torch.testing.assert_close(
sampled[field], flatten_image(full[field])[:, selected], rtol=0, atol=0
)
full_loss = sum(flatten_image(full[k])[:, selected].square().mean() for k in ("rgb", "depth"))
ray_loss = sum(sampled[k].square().mean() for k in ("rgb", "depth"))
full_grad = torch.autograd.grad(full_loss, tuple(model.parameters()))
ray_grad = torch.autograd.grad(ray_loss, tuple(model.parameters()))
for actual, expected in zip(ray_grad, full_grad, strict=True):
torch.testing.assert_close(actual, expected, rtol=2e-5, atol=1e-7)
assert actual.abs().sum() > 0
@pytest.mark.parametrize("with_depth", [True, False])
def test_optimizer_reduces_context_objective_selects_best_and_saves_detached_shared_state(
tmp_path, with_depth
):
batch, bounds = fixture_batch(depth=with_depth)
config = optimization.DirectOptimizationConfig(steps=30, batch_rays=64, checkpoint_interval=5)
output = tmp_path / "optimization"
state, summary = optimization.optimize_state(
batch, bounds, 4, 16, output, seed=91, config=config
)
checkpoints = summary["full_objective_checkpoints"]
best = min(checkpoints, key=lambda row: row["context_objective"])
assert summary["selected_step"] == best["step"]
assert summary["selected"]["context_objective"] == pytest.approx(
best["context_objective"], abs=1e-8
)
assert summary["selected"]["loss_rgb"] < summary["initial"]["loss_rgb"]
if with_depth:
assert summary["selected"]["loss_depth"] < summary["initial"]["loss_depth"]
else:
assert summary["selected"]["loss_depth"] is None
assert optimization.cache_rays(batch, bounds)["depth"] is None
assert summary["selected"]["context_objective"] < summary["initial"]["context_objective"]
assert (
summary["nonzero_gradient_observed"] and summary["all_steps_finite"] and summary["changed"]
)
assert summary["parameter_count"] == 4 * 4**3
assert summary["shared_supervision_frame_ids"] == [0, 3]
assert isinstance(state, SceneState)
assert state.features is None
for tensor in (state.density_logits, state.color, state.log_variance, state.bounds):
assert tensor.device.type == "cpu" and not tensor.requires_grad and tensor.grad_fn is None
assert torch.isfinite(tensor).all()
torch.testing.assert_close(state.log_variance, torch.full_like(state.log_variance, -3.0))
saved = torch.load(output / "state.pt", weights_only=False)
assert hash_scene_state(saved) == summary["state_hash"] == hash_scene_state(state)
before = hash_scene_state(state)
renderer = FixedMeasurementRenderer(n_samples=16)
initial = renderer(optimization.DirectState(bounds, 4).state(), batch.cameras)
final = renderer(state, batch.cameras)
assert not torch.equal(initial["rgb"], final["rgb"])
assert final["rgb"].shape[:2] == (1, 2)
assert hash_scene_state(state) == before # one unchanged state serves both cameras
lock = json.loads((output / "lock.json").read_text())
assert not lock["carrier_loaded"]
assert lock["has_depth"] == with_depth
with pytest.raises(FileExistsError):
optimization.optimize_state(batch, bounds, 4, 16, output, seed=91, config=config)
def test_zero_update_ties_select_initial_context_checkpoint(tmp_path):
batch, bounds = fixture_batch()
config = optimization.DirectOptimizationConfig(
steps=2, learning_rate=0, checkpoint_interval=1, batch_rays=8
)
_, summary = optimization.optimize_state(
batch, bounds, 4, 16, tmp_path / "ties", seed=1, config=config
)
assert summary["selected_step"] == 0
assert not summary["changed"]
assert len({r["context_objective"] for r in summary["full_objective_checkpoints"]}) == 1
@pytest.mark.parametrize("grid", [8, 16, 32])
def test_state_parameters_are_only_density_and_color(grid):
_, bounds = fixture_batch()
model = optimization.DirectState(bounds, grid)
assert set(dict(model.named_parameters())) == {"density", "color_logits"}
assert sum(p.numel() for p in model.parameters()) == 4 * grid**3
state = model.state()
assert state.density_logits is model.density
assert state.features is state.evidence is state.appearance is None
snapshot = clone_scene_state(state).to("cpu")
with torch.no_grad():
model.density.add_(1)
assert not torch.equal(snapshot.density_logits, model.density)
assert snapshot.density_logits.grad_fn is snapshot.color.grad_fn is None
@pytest.mark.parametrize("grid", [8, 16, 32])
def test_regularization_uses_physical_derivative_and_mean_density(grid):
bounds = torch.tensor([[-1.0, -2.0, -3.0], [1.0, 2.0, 3.0]])
model = optimization.DirectState(bounds, grid)
state = model.state()
# A color ramp with derivative 0.1 / metre along world x only.
centers = -1 + (torch.arange(grid) + 0.5) * (2 / grid)
state.color = (0.5 + 0.1 * centers).view(1, 1, 1, 1, grid).expand(1, 3, grid, grid, grid)
config = optimization.DirectOptimizationConfig()
expected = config.density_l2_weight * torch.nn.functional.softplus(torch.tensor(-2.0)).square()
expected += config.tv_weight * 0.1 / 6
torch.testing.assert_close(optimization.regularization(state, config), expected)
doubled_bounds = clone_scene_state(state)
doubled_bounds.bounds *= 2
no_l2 = replace(config, density_l2_weight=0)
torch.testing.assert_close(
optimization.regularization(doubled_bounds, no_l2),
optimization.regularization(state, no_l2) / 2,
)
def test_optimizer_has_no_query_or_checkpoint_api_and_no_carrier_import():
signature = inspect.signature(optimization.optimize_state)
assert set(signature.parameters) == {
"batch",
"bounds",
"grid",
"samples",
"output",
"seed",
"config",
}
assert all(p.kind != inspect.Parameter.VAR_KEYWORD for p in signature.parameters.values())
source = ast.parse(inspect.getsource(optimization))
imports = [node.module for node in ast.walk(source) if isinstance(node, ast.ImportFrom)]
assert not any(
"carrier" in name or "checkpoint" in name or "writer" in name for name in imports
)
calls = [node for node in ast.walk(source) if isinstance(node, ast.Call)]
assert not any(
isinstance(node.func, ast.Attribute) and node.func.attr == "load" for node in calls
)