feat(lib): add snap_endpoints to _postprocess
This commit is contained in:
@@ -29,3 +29,34 @@ def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
|
|||||||
smoothed[0] = x[0]
|
smoothed[0] = x[0]
|
||||||
smoothed[-1] = x[-1]
|
smoothed[-1] = x[-1]
|
||||||
return smoothed
|
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
|
||||||
|
|||||||
@@ -23,3 +23,24 @@ def test_gaussian_smooth_constant_unchanged() -> None:
|
|||||||
x = np.full(20, 7.5)
|
x = np.full(20, 7.5)
|
||||||
result = gaussian_smooth(x, sigma=1.0)
|
result = gaussian_smooth(x, sigma=1.0)
|
||||||
assert np.allclose(result, x, atol=1e-6)
|
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)
|
||||||
|
|||||||
Reference in New Issue
Block a user