diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py index 48f3a79..9124f38 100644 --- a/src/ai_mouse/_postprocess.py +++ b/src/ai_mouse/_postprocess.py @@ -29,3 +29,34 @@ def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray: smoothed[0] = x[0] smoothed[-1] = x[-1] return smoothed + + +def snap_endpoints( + forward: np.ndarray, + lateral: np.ndarray, + seq_len: int, + n_snap: int = 6, +) -> tuple[np.ndarray, np.ndarray]: + """Force first point to (0,0) and last point to (1,0) with quadratic ease. + + The last ``n_snap`` points are linearly interpolated towards (1, 0) + with quadratic easing, then the first/last points are pinned exactly. + + Args: + forward: (T,) forward coordinates (modified in place). + lateral: (T,) lateral coordinates (modified in place). + seq_len: length of forward/lateral. + n_snap: number of trailing points to ease (capped at seq_len//4). + + Returns: + ``(forward, lateral)`` after modification. + """ + n_snap = min(n_snap, seq_len // 4) + for i in range(n_snap): + alpha = ((i + 1) / n_snap) ** 2 + k = seq_len - n_snap + i + forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha + lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha + forward[0], lateral[0] = 0.0, 0.0 + forward[-1], lateral[-1] = 1.0, 0.0 + return forward, lateral diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py index a3d3198..4e3ba76 100644 --- a/tests/unit/test_postprocess.py +++ b/tests/unit/test_postprocess.py @@ -23,3 +23,24 @@ def test_gaussian_smooth_constant_unchanged() -> None: x = np.full(20, 7.5) result = gaussian_smooth(x, sigma=1.0) assert np.allclose(result, x, atol=1e-6) + + +from ai_mouse._postprocess import snap_endpoints + + +def test_snap_endpoints_pins_first_and_last() -> None: + forward = np.linspace(0.1, 0.9, 16) + lateral = np.full(16, 0.5) + f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16) + assert f[0] == 0.0 + assert l[0] == 0.0 + assert f[-1] == 1.0 + assert l[-1] == 0.0 + + +def test_snap_endpoints_preserves_middle() -> None: + forward = np.linspace(0.0, 1.0, 16) + lateral = np.zeros(16) + f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4) + # Points before the last n_snap should be unchanged + assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6)