diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py index 19f8bcd..9072159 100644 --- a/src/ai_mouse/_postprocess.py +++ b/src/ai_mouse/_postprocess.py @@ -237,3 +237,36 @@ def soften_forward( over = out > 1.0 out[over] = 1.0 + overshoot_span * np.tanh((out[over] - 1.0) / overshoot_span) return out + + +def warp_endpoints( + forward: np.ndarray, + lateral: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Warp the whole curve so endpoints land exactly on (0,0) and (1,0). + + Computes the residual of the first point vs (0, 0) and the last point + vs (1, 0), then subtracts each residual weighted by a smoothstep that + is 1 at its own end and 0 at the opposite end. Unlike tail-dragging, + this preserves the trajectory's local shape and approach direction. + + Args: + forward: (T,) forward coordinates. + lateral: (T,) lateral coordinates. + + Returns: + ``(forward, lateral)`` new arrays, endpoints pinned exactly. + """ + t = np.linspace(0.0, 1.0, len(forward)) + w_end = t * t * (3.0 - 2.0 * t) # smoothstep: 0 at start → 1 at end + w_start = 1.0 - w_end # mirrored + + res_f0, res_l0 = forward[0] - 0.0, lateral[0] - 0.0 + res_f1, res_l1 = forward[-1] - 1.0, lateral[-1] - 0.0 + + fo = forward - w_start * res_f0 - w_end * res_f1 + lo = lateral - w_start * res_l0 - w_end * res_l1 + + fo[0], lo[0] = 0.0, 0.0 + fo[-1], lo[-1] = 1.0, 0.0 + return fo, lo diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py index e2bcd2e..ffd68ba 100644 --- a/tests/unit/test_postprocess.py +++ b/tests/unit/test_postprocess.py @@ -206,3 +206,45 @@ def test_damp_start_short_input_safe() -> None: lat = np.array([0.0, 0.5, 0.3]) out = damp_start(lat, n=4) # n capped to len//4 = 0 → no-op assert np.array_equal(out, lat) + + +from ai_mouse._postprocess import warp_endpoints + + +def test_warp_endpoints_exact_pin() -> None: + f = np.linspace(0.05, 1.10, 32) + l = np.linspace(0.03, -0.07, 32) + fo, lo = warp_endpoints(f, l) + assert fo[0] == 0.0 and lo[0] == 0.0 + assert fo[-1] == 1.0 and lo[-1] == 0.0 + + +def test_warp_endpoints_identity_when_already_pinned() -> None: + f = np.linspace(0.0, 1.0, 32) + l = np.sin(np.linspace(0, np.pi, 32)) * 0.1 + l[0] = l[-1] = 0.0 + fo, lo = warp_endpoints(f.copy(), l.copy()) + assert np.allclose(fo, f, atol=1e-12) + assert np.allclose(lo, l, atol=1e-12) + + +def test_warp_endpoints_preserves_smoothness() -> None: + # Correcting a smooth curve must not introduce sharp local bends: + # the warp adds a smoothstep-weighted offset, so the second + # difference (discrete curvature proxy) stays small. + f = np.linspace(0.02, 1.08, 32) + l = np.full(32, 0.05) + fo, lo = warp_endpoints(f, l) + assert np.abs(np.diff(lo, 2)).max() < 0.01 + assert np.abs(np.diff(fo, 2)).max() < 0.01 + + +def test_warp_endpoints_correction_local_to_each_end() -> None: + # A start-only residual should barely move the last quarter. + # (smoothstep weight at i=24/31 is ~0.13, so 0.08 residual leaves + # ~0.010 there — threshold 0.02 gives margin without losing meaning) + f = np.linspace(0.0, 1.0, 32) + 0.0 + l = np.zeros(32) + f[0] = 0.08 # start residual only + fo, _ = warp_endpoints(f, l) + assert np.abs(fo[24:] - f[24:]).max() < 0.02