From 6cfcb6d1a4f3ab84b506de56c9efab2b5266f009 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 01:07:58 +0800
Subject: [PATCH] feat(lib): add smooth_start, enforce_forward_monotonic

---
 src/ai_mouse/_postprocess.py   | 30 ++++++++++++++++++++++++++++++
 tests/unit/test_postprocess.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py
index 9124f38..ac2b6e3 100644
--- a/src/ai_mouse/_postprocess.py
+++ b/src/ai_mouse/_postprocess.py
@@ -60,3 +60,33 @@ def snap_endpoints(
     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
diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py
index 4e3ba76..0ec1439 100644
--- a/tests/unit/test_postprocess.py
+++ b/tests/unit/test_postprocess.py
@@ -44,3 +44,31 @@ def test_snap_endpoints_preserves_middle() -> None:
     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