feat: rework mouse post-processing pipeline (soft monotonic, global endpoint warp)

Golden mouse baselines temporarily failing; re-captured in follow-up commit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

src/ai_mouse/_postprocess.py

@@ -31,67 +31,6 @@ def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
     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
-
-
-def smooth_start(
-    forward: np.ndarray,
-    lateral: np.ndarray,
-    n: int = 4,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Dampen lateral oscillation in the first ``n`` points.
-
-    Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is
-    forced non-decreasing locally; lateral is linearly damped towards 0.
-    """
-    n_start_fix = min(n, len(forward) // 4)
-    for i in range(1, n_start_fix + 1):
-        blend = i / (n_start_fix + 1)
-        forward[i] = max(forward[i], forward[i - 1])
-        lateral[i] = lateral[i] * blend
-    return forward, lateral
-
-
-def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray:
-    """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints."""
-    seq_len = len(forward)
-    for i in range(1, seq_len - 1):
-        if forward[i] < forward[i - 1]:
-            forward[i] = forward[i - 1] + 0.001
-    forward = np.clip(forward, 0.0, 1.0)
-    forward[0] = 0.0
-    forward[-1] = 1.0
-    return forward
-
-
 def damp_start(lateral: np.ndarray, n: int = 4) -> np.ndarray:
     """Dampen lateral oscillation over the first ``n`` points, continuously.
 

src/ai_mouse/mouse.py

@@ -14,13 +14,13 @@ from ai_mouse._assets import resolve
 from ai_mouse._coord import decode_trajectory
 from ai_mouse._postprocess import (
     build_timestamps,
-    enforce_forward_monotonic,
+    damp_start,
     gaussian_smooth,
     resample_arc,
     sample_duration,
-    smooth_start,
+    soften_forward,
-    snap_endpoints,
     truncnorm_sample,
+    warp_endpoints,
 )
 from ai_mouse.errors import GenerationError, ModelLoadError
 
@@ -103,10 +103,11 @@ class MouseModel:
         lateral = x[0, :, 1].copy()
         log_dt = x[0, :, 2].copy()
 
-        forward, lateral = snap_endpoints(forward, lateral, self._seq_len)
+        forward = soften_forward(forward)
-        forward, lateral = smooth_start(forward, lateral)
+        lateral = damp_start(lateral)
-        forward = enforce_forward_monotonic(forward)
+        forward = gaussian_smooth(forward, sigma=1.0)
         lateral = gaussian_smooth(lateral, sigma=1.0)
+        forward, lateral = warp_endpoints(forward, lateral)
 
         log_dt = np.clip(log_dt, 0.0, 5.0)
         log_dt[0] = 0.0

tests/unit/test_postprocess.py

@@ -25,55 +25,6 @@ def test_gaussian_smooth_constant_unchanged() -> None:
     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)
-
-
-from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start
-
-
-def test_smooth_start_dampens_lateral() -> None:
-    forward = np.linspace(0, 1, 16)
-    lateral = np.full(16, 1.0)
-    forward[0] = lateral[0] = 0.0  # invariant: snap already done
-    _, l = smooth_start(forward.copy(), lateral.copy(), n=4)
-    # Lateral at points 1-4 should be < original (dampened)
-    assert l[1] < 1.0
-    assert l[4] < 1.0
-    # Lateral at point 5+ unchanged
-    assert l[5] == 1.0
-
-
-def test_enforce_forward_monotonic_repairs_inversions() -> None:
-    f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0])
-    out = enforce_forward_monotonic(f.copy())
-    assert np.all(np.diff(out) > 0), out
-
-
-def test_enforce_forward_monotonic_clips_to_unit_interval() -> None:
-    f = np.array([-0.1, 0.5, 1.2])
-    out = enforce_forward_monotonic(f.copy())
-    assert out[0] == 0.0
-    assert out[-1] == 1.0
-
-
 from ai_mouse._postprocess import build_timestamps, resample_arc