Additional pipeline tests

.github/workflows/ci.yml

@@ -4,7 +4,6 @@ on:
   push:
     branches: [main]
   pull_request:
-    branches: [main]
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}

tests/reanchor_reference.py

@@ -0,0 +1,12 @@
+"""Reference implementation of the LoGeR re-anchor transform."""
+
+import numpy as np
+
+
+def single_shot_reanchor_world(poses, world_points):
+    # One homogeneous matmul over all points. The gold standard the block loop
+    # must reproduce bit-for-bit.
+    reference_inv = np.linalg.inv(poses.astype(np.float64)[0])
+    flat = world_points.astype(np.float64).reshape(-1, 3)
+    homog = np.concatenate([flat, np.ones((flat.shape[0], 1), dtype=np.float64)], axis=1)
+    return (homog @ reference_inv.T)[:, :3].reshape(world_points.shape).astype(world_points.dtype)

tests/test_device.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import torch
+
+from deepreefmap.device import (
+    autocast_context,
+    disable_torch_compile_without_triton,
+    get_autocast_dtype,
+    release_device_memory,
+    resolve_device,
+)
+
+
+def test_resolve_device_prefers_cuda() -> None:
+    with patch("torch.cuda.is_available", return_value=True):
+        assert resolve_device().type == "cuda"
+
+
+def test_resolve_device_prefers_cuda_over_mps() -> None:
+    with (
+        patch("torch.cuda.is_available", return_value=True),
+        patch("torch.backends.mps.is_available", return_value=True),
+    ):
+        assert resolve_device().type == "cuda"
+
+
+def test_resolve_device_prefers_mps_over_cpu() -> None:
+    with (
+        patch("torch.cuda.is_available", return_value=False),
+        patch("torch.backends.mps.is_available", return_value=True),
+    ):
+        assert resolve_device().type == "mps"
+
+
+def test_resolve_device_falls_back_to_cpu() -> None:
+    with (
+        patch("torch.cuda.is_available", return_value=False),
+        patch("torch.backends.mps.is_available", return_value=False),
+    ):
+        assert resolve_device().type == "cpu"
+
+
+def test_get_autocast_dtype_mps() -> None:
+    assert get_autocast_dtype(torch.device("mps")) == torch.float16
+
+
+def test_get_autocast_dtype_cpu() -> None:
+    assert get_autocast_dtype(torch.device("cpu")) == torch.bfloat16
+
+
+def test_get_autocast_dtype_safe_on_capability_error() -> None:
+    with patch("torch.cuda.get_device_capability", side_effect=RuntimeError("no device")):
+        assert get_autocast_dtype(torch.device("cuda")) == torch.float16
+
+
+def test_get_autocast_dtype_cuda_bf16_needs_flash() -> None:
+    with patch("torch.cuda.get_device_capability", return_value=(12, 0)):
+        with patch("deepreefmap.device._flash_sdpa_works", return_value=True):
+            assert get_autocast_dtype(torch.device("cuda")) == torch.bfloat16
+        with patch("deepreefmap.device._flash_sdpa_works", return_value=False):
+            assert get_autocast_dtype(torch.device("cuda")) == torch.float16
+
+
+def test_get_autocast_dtype_old_gpu_skips_flash_probe() -> None:
+    # The probe answers yes, so only the capability gate can force float16.
+    with (
+        patch("torch.cuda.get_device_capability", return_value=(7, 5)),
+        patch("deepreefmap.device._flash_sdpa_works", return_value=True),
+    ):
+        assert get_autocast_dtype(torch.device("cuda")) == torch.float16
+
+
+def test_disable_torch_compile_without_triton_disables_dynamo() -> None:
+    original = torch._dynamo.config.disable
+    try:
+        with patch("deepreefmap.device.importlib.util.find_spec", return_value=None):
+            torch._dynamo.config.disable = False
+            disable_torch_compile_without_triton()
+            assert torch._dynamo.config.disable is True
+    finally:
+        torch._dynamo.config.disable = original
+
+
+def test_disable_torch_compile_with_triton_leaves_dynamo_alone() -> None:
+    original = torch._dynamo.config.disable
+    try:
+        with patch("deepreefmap.device.importlib.util.find_spec", return_value=object()):
+            torch._dynamo.config.disable = False
+            disable_torch_compile_without_triton()
+            assert torch._dynamo.config.disable is False
+    finally:
+        torch._dynamo.config.disable = original
+
+
+def test_autocast_context_returns_context_manager() -> None:
+    ctx = autocast_context(torch.device("cpu"))
+    assert hasattr(ctx, "__enter__") and hasattr(ctx, "__exit__")
+
+
+def test_release_device_memory_cpu_no_error() -> None:
+    release_device_memory(torch.device("cpu"))

tests/test_loger_progress.py

@@ -0,0 +1,38 @@
+"""LoGeR progress reporting: window counting and per-block re-anchor progress."""
+
+import numpy as np
+from reanchor_reference import single_shot_reanchor_world
+
+from deepreefmap.mapping.loger_backend import _count_windows, _reanchor_to_first_camera
+
+
+def test_count_windows_matches_pi3_sliding_scheme() -> None:
+    # Whole sequence fits in one window → single pass, no countdown.
+    assert _count_windows(10, window_size=32, overlap_size=3) == 1
+    assert _count_windows(10, window_size=0, overlap_size=3) == 1
+    # step = window - overlap = 29; windows start at 0, 29, 58, ... until N.
+    assert _count_windows(100, window_size=32, overlap_size=3) == 4
+    assert _count_windows(64, window_size=32, overlap_size=3) == 3
+
+
+def test_reanchor_reports_monotonic_block_progress(monkeypatch) -> None:
+    import deepreefmap.mapping.loger_backend as lb
+
+    rng = np.random.default_rng(5)
+    n, h, w = 4, 5, 7  # 140 points
+    poses = np.tile(np.eye(4, dtype=np.float32), (n, 1, 1))
+    poses[0, :3, 3] = (1.0, 0.5, -0.25)  # non-trivial transform, not a no-op check
+    world = (rng.random((n, h, w, 3), dtype=np.float64).astype(np.float32) * 2.0 - 1.0)
+    reference = single_shot_reanchor_world(poses, world)
+
+    calls: list[tuple[int, int, str]] = []
+    monkeypatch.setattr(lb, "_REANCHOR_POINT_BLOCK", 40)
+    _, rebased_world = _reanchor_to_first_camera(
+        poses, world, lambda cur, tot, msg: calls.append((cur, tot, msg))
+    )
+    # Progress reporting must not perturb the transform.
+    assert np.array_equal(rebased_world, reference)
+    # 140 points in blocks of 40 -> reports at 40, 80, 120, 140, all against 140.
+    assert [c[0] for c in calls] == [40, 80, 120, 140]
+    assert all(tot == 140 for _, tot, _ in calls)
+    assert all(msg == "Aligning poses to world frame" for _, _, msg in calls)

tests/test_loger_reanchor.py

@@ -1,15 +1,8 @@
 import numpy as np
+from reanchor_reference import single_shot_reanchor_world
 
 from deepreefmap.mapping.loger_backend import LoGeRBackend, _reanchor_to_first_camera
 
-def _single_shot_reanchor_world(poses, world_points):
-    # The pre-chunking transform: one homogeneous matmul over all points. The
-    # gold standard the block loop must reproduce bit-for-bit.
-    reference_inv = np.linalg.inv(poses.astype(np.float64)[0])
-    flat = world_points.astype(np.float64).reshape(-1, 3)
-    homog = np.concatenate([flat, np.ones((flat.shape[0], 1), dtype=np.float64)], axis=1)
-    return (homog @ reference_inv.T)[:, :3].reshape(world_points.shape).astype(world_points.dtype)
-
 
 def test_reanchor_chunking_is_bitwise_identical_across_block_sizes(monkeypatch) -> None:
     import deepreefmap.mapping.loger_backend as lb
@@ -21,7 +14,7 @@ def test_reanchor_chunking_is_bitwise_identical_across_block_sizes(monkeypatch)
     poses[1, :3, :3] = np.linalg.qr(rng.random((3, 3)))[0].astype(np.float32)
     world = (rng.random((n, h, w, 3), dtype=np.float64).astype(np.float32) * 2.0 - 1.0)
 
-    reference = _single_shot_reanchor_world(poses, world)
+    reference = single_shot_reanchor_world(poses, world)
     # Block sizes that split the n*h*w=140 points at every awkward boundary.
     # Re-anchoring consumes its input, so each call gets a fresh copy.
     for block in (1, 2, 3, 13, 139, 140, 10_000):
@@ -90,7 +83,7 @@ def test_sequence_result_keeps_unanchored_local_points_when_pi3_omits_them() ->
     result = backend.process_sequence([0, 1], images)
 
     assert np.array_equal(result.local_points, world[0])
-    expected_world = _single_shot_reanchor_world(poses[0], world[0])
+    expected_world = single_shot_reanchor_world(poses[0], world[0])
     assert np.array_equal(result.world_points, expected_world)
     assert result.depth_maps.dtype == np.float32
     assert result.world_points.dtype == np.float32

tests/test_mapping_progress.py

@@ -0,0 +1,69 @@
+"""The base process_sequence loop reports per-frame progress to a callback."""
+
+from __future__ import annotations
+
+import threading
+
+import numpy as np
+import pytest
+
+from deepreefmap.mapping.base import FrameEstimate, MappingBackend
+from deepreefmap.pipeline.artifacts import ReconstructionCancelled
+
+
+class _StubBackend(MappingBackend):
+    name = "stub"
+    default_window_size = 1
+
+    def initialize(self, image_size, intrinsics) -> None:
+        del image_size, intrinsics
+
+    def process_frame(self, frame_index: int, image_rgb: np.ndarray) -> FrameEstimate:
+        del image_rgb
+        return FrameEstimate(
+            frame_index=frame_index,
+            depth=np.ones((2, 2), dtype=np.float32),
+            pose_w_c=np.eye(4, dtype=np.float32),
+            intrinsics=np.eye(3, dtype=np.float32),
+        )
+
+
+def _images(n: int) -> list[np.ndarray]:
+    return [np.zeros((2, 2, 3), dtype=np.uint8) for _ in range(n)]
+
+
+def test_progress_callback_fires_once_per_frame() -> None:
+    calls: list[tuple[int, int, str]] = []
+    backend = _StubBackend()
+    backend.process_sequence(
+        [10, 11, 12],
+        _images(3),
+        progress_callback=lambda cur, tot, msg: calls.append((cur, tot, msg)),
+    )
+    assert [c[:2] for c in calls] == [(1, 3), (2, 3), (3, 3)]
+    assert all(msg for _, _, msg in calls)
+
+
+def test_process_sequence_without_callback_still_runs() -> None:
+    backend = _StubBackend()
+    result = backend.process_sequence([0, 1], _images(2))
+    assert result.depth_maps.shape[0] == 2
+
+
+def test_cancel_event_stops_mid_sequence() -> None:
+    cancel = threading.Event()
+    calls: list[int] = []
+
+    def record_and_cancel(cur: int, tot: int, msg: str) -> None:
+        calls.append(cur)
+        if cur == 2:
+            cancel.set()
+
+    with pytest.raises(ReconstructionCancelled):
+        _StubBackend().process_sequence([0, 1, 2, 3], _images(4), cancel_event=cancel, progress_callback=record_and_cancel)
+    assert calls == [1, 2]
+
+
+def test_empty_sequence_raises() -> None:
+    with pytest.raises(RuntimeError, match="empty mapping sequence"):
+        _StubBackend().process_sequence([], [])