From 6d848dc23dce5fb481edf0d9d43124327a2c5702 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 01:06:42 +0800
Subject: [PATCH] feat(lib): add gaussian_smooth to _postprocess

---
 src/ai_mouse/_postprocess.py   | 31 +++++++++++++++++++++++++++++++
 tests/unit/test_postprocess.py | 25 +++++++++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 src/ai_mouse/_postprocess.py
 create mode 100644 tests/unit/test_postprocess.py

diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py
new file mode 100644
index 0000000..48f3a79
--- /dev/null
+++ b/src/ai_mouse/_postprocess.py
@@ -0,0 +1,31 @@
+"""Pure-numpy post-processing primitives for trajectory generation.
+
+All functions are pure (no I/O, no global state) and accept an explicit
+:class:`numpy.random.Generator` when randomness is involved.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+
+def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
+    """5-tap gaussian smoothing along a 1-D array; endpoints preserved.
+
+    Args:
+        x: 1-D input array.
+        sigma: gaussian std. Default 1.0 gives weights approximately
+               [0.054, 0.244, 0.403, 0.244, 0.054].
+
+    Returns:
+        Smoothed array of the same shape. ``x[0]`` and ``x[-1]`` unchanged.
+        If ``len(x) < 5`` returns a copy of ``x`` (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()
+    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
diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py
new file mode 100644
index 0000000..a3d3198
--- /dev/null
+++ b/tests/unit/test_postprocess.py
@@ -0,0 +1,25 @@
+"""Tests for trajectory post-processing primitives."""
+from __future__ import annotations
+
+import numpy as np
+
+from ai_mouse._postprocess import gaussian_smooth
+
+
+def test_gaussian_smooth_preserves_endpoints() -> None:
+    x = np.array([1.0, 5.0, 3.0, 8.0, 2.0, 6.0, 4.0])
+    result = gaussian_smooth(x, sigma=1.0)
+    assert result[0] == 1.0
+    assert result[-1] == 4.0
+
+
+def test_gaussian_smooth_short_input_unchanged() -> None:
+    x = np.array([1.0, 2.0, 3.0])
+    result = gaussian_smooth(x, sigma=1.0)
+    assert np.array_equal(result, x)
+
+
+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)