tests/test_velocity.py

"""Synthetic, data-free tests for fpgm.geometry.velocity."""

from __future__ import annotations

import numpy as np

from fpgm.config import VelocityConfig
from fpgm.geometry.velocity import VelocityEstimator
from fpgm.types import Track3D


def _linear_track(
    velocities: np.ndarray, p0: np.ndarray, n_frames: int, dt: float = 0.1
) -> Track3D:
    """Build a Track3D where point q moves at constant ``velocities[q]`` from ``p0[q]``."""
    q = velocities.shape[0]
    t = np.arange(n_frames) * dt
    xyz = p0[None, :, :] + velocities[None, :, :] * t[:, None, None]
    return Track3D(
        point_id=np.arange(q, dtype=np.int32),
        timestamps=t,
        xyz_world=xyz,
        valid=np.ones((n_frames, q), dtype=bool),
        confidence=np.ones((n_frames, q), dtype=np.float32),
    )


class TestConstantVelocity:
    def test_recovers_exact_velocity(self):
        v_true = np.array([[1.0, -0.5, 2.0]])
        p0 = np.array([[0.1, 0.2, 0.3]])
        track = _linear_track(v_true, p0, n_frames=20)

        estimate = VelocityEstimator().estimate(track, VelocityConfig())

        assert np.allclose(estimate.linear_velocity[:, 0, :], v_true[0], atol=1e-6)
        assert np.allclose(estimate.object_linear_velocity, v_true[0], atol=1e-6)
        assert np.allclose(estimate.object_speed, np.linalg.norm(v_true[0]), atol=1e-6)
        assert estimate.frame == "world"


class TestGapHandling:
    def test_short_gap_is_bridged(self):
        v_true = np.array([[0.5, 0.0, -0.2]])
        p0 = np.array([[0.0, 0.0, 1.0]])
        n_frames = 15
        track = _linear_track(v_true, p0, n_frames=n_frames)

        # A 2-frame gap, within max_gap_frames=3 (the default): should be bridged.
        track.valid[6:8, 0] = False
        track.xyz_world[6:8, 0, :] = np.nan

        cfg = VelocityConfig()
        estimate = VelocityEstimator().estimate(track, cfg)

        assert not np.any(np.isnan(estimate.linear_velocity[:, 0, :]))
        assert np.allclose(estimate.linear_velocity[:, 0, :], v_true[0], atol=1e-4)

    def test_long_gap_yields_nan_across_gap_and_correct_velocity_on_both_sides(self):
        v_true = np.array([[0.5, 0.0, -0.2]])
        p0 = np.array([[0.0, 0.0, 1.0]])
        n_frames = 15
        track = _linear_track(v_true, p0, n_frames=n_frames)

        cfg = VelocityConfig()  # max_gap_frames=3
        # A 5-frame gap: longer than max_gap_frames, must NOT be bridged.
        gap = slice(5, 10)
        track.valid[gap, 0] = False
        track.xyz_world[gap, 0, :] = np.nan

        estimate = VelocityEstimator().estimate(track, cfg)

        assert np.all(np.isnan(estimate.linear_velocity[gap, 0, :]))
        assert np.allclose(estimate.linear_velocity[0:5, 0, :], v_true[0], atol=1e-4)
        assert np.allclose(estimate.linear_velocity[10:15, 0, :], v_true[0], atol=1e-4)


class TestRobustAggregation:
    def test_single_outlier_does_not_move_median_object_velocity(self):
        v0 = np.array([1.0, -0.5, 0.3])
        # Four points move at v0 plus tiny, distinct jitter; one point is a gross
        # outlier (velocity offset by +10 m/s in every axis).
        velocities = np.array(
            [
                v0 + [0.01, 0.0, 0.0],
                v0 + [-0.01, 0.0, 0.0],
                v0 + [0.0, 0.01, 0.0],
                v0 + [0.0, -0.01, 0.0],
                v0 + [10.0, 10.0, 10.0],
            ]
        )
        outlier_idx = 4
        p0 = np.zeros((5, 3))
        track = _linear_track(velocities, p0, n_frames=20)

        estimate = VelocityEstimator().estimate(track, VelocityConfig())

        assert np.allclose(estimate.object_linear_velocity, v0, atol=0.02)
        mid = 10
        assert not estimate.inlier_mask[mid, outlier_idx]
        assert np.all(estimate.inlier_mask[mid, :outlier_idx])