tests/test_resample_trajectory.py

"""Regression tests for ``scripts/resample_trajectory.py``.

These exist because of one specific bug that shipped a visibly broken render
while every self-check passed. Object transforms in this pipeline are **not**
bare rotations: the planner bakes the asset's own scale into the 3x3 block
(the brick carries a uniform factor of ~0.05). Resampling them through
``scipy.spatial.transform.Rotation`` orthonormalises that scale away, so the
object renders ~20x too large -- and an "is the 3x3 block a valid rotation?"
assertion *passes*, because destroying the scale is exactly what makes it one.
The test below therefore checks the property that actually matters (scale is
preserved) rather than the property that is easy to assert.
"""

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path

import numpy as np
import pytest
from scipy.spatial.transform import Rotation

REPO_ROOT = Path(__file__).resolve().parents[1]


def _load_script():
    path = REPO_ROOT / "scripts" / "resample_trajectory.py"
    spec = importlib.util.spec_from_file_location("resample_trajectory", path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = mod  # dataclass/typing introspection needs this
    spec.loader.exec_module(mod)
    return mod


rt = _load_script()


def _scaled_pose_track(n: int, scale: np.ndarray) -> np.ndarray:
    """(n, 1, 4, 4) transforms that rotate and translate while carrying `scale`."""
    angles = np.linspace(0.0, 0.9, n)
    poses = np.zeros((n, 1, 4, 4))
    poses[:, 0, 3, 3] = 1.0
    for i, a in enumerate(angles):
        R = Rotation.from_rotvec([0.3 * a, -0.5 * a, 0.8 * a]).as_matrix()
        poses[i, 0, :3, :3] = R * scale[None, :]
        poses[i, 0, :3, 3] = [0.1 * a, 0.2 - 0.05 * a, 0.4 + 0.3 * a]
    return poses


@pytest.mark.parametrize(
    "scale",
    [
        np.array([0.050199, 0.050199, 0.050199]),  # the real brick's factor
        np.array([1.0, 1.0, 1.0]),                 # already metric
        np.array([0.02, 0.05, 0.31]),              # non-uniform
    ],
)
def test_resample_preserves_object_transform_scale(scale):
    src = _scaled_pose_track(40, scale)
    s_src = np.arange(40, dtype=float)
    s_dst = np.linspace(0.0, 39.0, 57)

    out = rt._resample_poses(src, s_src, s_dst)

    got = np.linalg.norm(out[:, 0, :3, :3], axis=1)
    assert np.allclose(got, scale[None, :], atol=1e-12), (
        f"scale not preserved: expected {scale}, got column norms in "
        f"[{got.min(axis=0)}, {got.max(axis=0)}]"
    )


def test_resample_preserves_rotation_and_endpoints():
    scale = np.array([0.050199, 0.050199, 0.050199])
    src = _scaled_pose_track(40, scale)
    s_src = np.arange(40, dtype=float)
    s_dst = np.linspace(0.0, 39.0, 57)

    out = rt._resample_poses(src, s_src, s_dst)

    # Endpoints land exactly on source samples, so they must round-trip bitwise-close.
    assert np.allclose(out[0, 0], src[0, 0], atol=1e-12)
    assert np.allclose(out[-1, 0], src[-1, 0], atol=1e-12)

    # The rotation factor must stay a rotation despite carrying the scale.
    R = out[:, 0, :3, :3] / scale[None, None, :]
    assert np.allclose(R @ np.swapaxes(R, 1, 2), np.eye(3), atol=1e-9)
    assert np.allclose(np.linalg.det(R), 1.0, atol=1e-9)


def test_resample_rejects_degenerate_transform():
    src = _scaled_pose_track(10, np.array([0.05, 0.05, 0.05]))
    src[3, 0, :3, 2] = 0.0  # collapse one axis
    with pytest.raises(SystemExit, match="degenerate"):
        rt._resample_poses(src, np.arange(10, dtype=float), np.linspace(0, 9, 13))


def test_resample_interpolates_a_translating_object_linearly():
    """A constant-velocity object must stay on its line after resampling."""
    n = 30
    poses = np.zeros((n, 1, 4, 4))
    poses[:, 0, 3, 3] = 1.0
    poses[:, 0, :3, :3] = np.eye(3) * 0.05
    t = np.linspace(0.0, 1.0, n)
    poses[:, 0, :3, 3] = np.stack([t, 2 * t, 0.5 - t], axis=1)

    out = rt._resample_poses(poses, np.arange(n, dtype=float), np.linspace(0, n - 1, 47))
    p = out[:, 0, :3, 3]
    assert np.allclose(p[:, 1], 2 * p[:, 0], atol=1e-12)
    assert np.allclose(p[:, 2], 0.5 - p[:, 0], atol=1e-12)


def test_chord_error_is_zero_for_a_linear_signal():
    """The validation metric must not manufacture error where there is none."""
    s = np.arange(21, dtype=float)
    values = np.stack([3.0 * s - 1.0, -0.5 * s], axis=1)
    assert rt._chord_error(values, s) == pytest.approx(0.0, abs=1e-12)


def test_chord_error_detects_curvature():
    s = np.arange(21, dtype=float)
    values = (s ** 2)[:, None]
    assert rt._chord_error(values, s) > 0.5