feat(lib): add smooth_start, enforce_forward_monotonic
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user