From 3fb4d3a8c0363fb140567a52b98bd0134edabec0 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Sun, 10 May 2026 13:27:15 +0800
Subject: [PATCH] feat(generator): add 5-point gaussian smoothing on lateral

---
 ai_mouse/generator.py   | 33 +++++++++++++++++++++++++++++++++
 tests/test_generator.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/ai_mouse/generator.py b/ai_mouse/generator.py
index 56e63c1..6b05055 100644
--- a/ai_mouse/generator.py
+++ b/ai_mouse/generator.py
@@ -10,6 +10,7 @@ Pipeline:
    a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
    b. Smooth start: dampen lateral near start (first 4 points).
    c. Enforce forward monotonicity (prevent x-axis jitter).
+   d. 5-point gaussian smooth on lateral (preserve endpoints).
 6. Temporal post-processing:
    a. Clip log_dt to [0, 5] to prevent exponential explosion.
    (speed profile and median±1.1 hard clip removed in 2026-05 refactor —
@@ -73,6 +74,35 @@ def _sample_duration(duration_dist: dict, dist: float) -> float:
     return float(np.exp(np.random.normal(mu_log, sigma_log)))
 
 
+# ---------------------------------------------------------------------------
+# Smoothing helper
+# ---------------------------------------------------------------------------
+
+
+def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
+    """5-point gaussian smoothing along a 1-D array, preserving endpoints.
+
+    Args:
+        x: 1-D input array.
+        sigma: Gaussian std (px); larger = more smoothing. Default 1.0 gives
+               weights ≈ [0.054, 0.244, 0.403, 0.244, 0.054].
+
+    Returns:
+        Smoothed array of the same shape. x[0] and x[-1] are unchanged.
+        If len(x) < 5, returns x unchanged (kernel won't fit).
+    """
+    if len(x) < 5:
+        return x.copy()
+    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
+    kernel /= kernel.sum()
+    # Pad with edge values to avoid boundary artifacts, then slice back
+    padded = np.pad(x, pad_width=2, mode="edge")
+    smoothed = np.convolve(padded, kernel, mode="valid")
+    smoothed[0] = x[0]
+    smoothed[-1] = x[-1]
+    return smoothed
+
+
 # ---------------------------------------------------------------------------
 # Main generate function
 # ---------------------------------------------------------------------------
@@ -245,6 +275,9 @@ def generate(
     forward[0] = 0.0
     forward[-1] = 1.0
 
+    # Lateral 5-point gaussian smoothing (endpoints preserved)
+    lateral = _gaussian_smooth(lateral, sigma=1.0)
+
     # ------------------------------------------------------------------
     # Temporal post-processing (log_dt)
     # ------------------------------------------------------------------
diff --git a/tests/test_generator.py b/tests/test_generator.py
index 4f3ce06..4bbf09d 100644
--- a/tests/test_generator.py
+++ b/tests/test_generator.py
@@ -124,3 +124,35 @@ class TestPostProcessing:
             if len(values) > 1:
                 return  # at least one position has variation — pass
         pytest.fail("All 5 Δt sequences are identical at every position — diversity collapsed")
+
+
+class TestGaussianSmooth:
+    def test_endpoints_preserved(self):
+        from ai_mouse.generator import _gaussian_smooth
+        x = np.array([1.0, 5.0, 3.0, 7.0, 2.0], dtype=np.float64)
+        smoothed = _gaussian_smooth(x, sigma=1.0)
+        assert smoothed[0] == 1.0
+        assert smoothed[-1] == 2.0
+
+    def test_smooths_high_frequency(self):
+        """A high-frequency square wave should have reduced amplitude after smoothing."""
+        from ai_mouse.generator import _gaussian_smooth
+        x = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=np.float64)
+        smoothed = _gaussian_smooth(x, sigma=1.0)
+        # Interior amplitude should be reduced
+        interior_orig = x[2:-2]
+        interior_smooth = smoothed[2:-2]
+        assert interior_smooth.std() < interior_orig.std()
+
+    def test_constant_signal_unchanged(self):
+        from ai_mouse.generator import _gaussian_smooth
+        x = np.full(20, 0.5, dtype=np.float64)
+        smoothed = _gaussian_smooth(x, sigma=1.0)
+        np.testing.assert_allclose(smoothed, x, rtol=1e-6)
+
+    def test_short_array_returns_unchanged(self):
+        """Arrays shorter than the kernel are returned unchanged."""
+        from ai_mouse.generator import _gaussian_smooth
+        x = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+        smoothed = _gaussian_smooth(x, sigma=1.0)
+        np.testing.assert_allclose(smoothed, x, rtol=1e-6)