dog
441e6f3dfe
Golden mouse baselines temporarily failing; re-captured in follow-up commit. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
202 lines
6.5 KiB
Python
202 lines
6.5 KiB
Python
"""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)
|
|
|
|
|
|
from ai_mouse._postprocess import build_timestamps, resample_arc
|
|
|
|
|
|
def test_resample_arc_identity_when_same_length() -> None:
|
|
pts = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0], [3.0, 1.0]])
|
|
out = resample_arc(pts, 4)
|
|
assert np.allclose(out, pts, atol=1e-6)
|
|
|
|
|
|
def test_resample_arc_changes_length() -> None:
|
|
pts = np.array([[float(i), 0.0] for i in range(10)])
|
|
out = resample_arc(pts, 5)
|
|
assert out.shape == (5, 2)
|
|
# Endpoints preserved
|
|
assert np.allclose(out[0], pts[0])
|
|
assert np.allclose(out[-1], pts[-1])
|
|
|
|
|
|
def test_build_timestamps_strictly_increasing() -> None:
|
|
log_dt = np.array([0.0, 2.0, 2.5, 3.0, 2.0])
|
|
ts = build_timestamps(log_dt, total_duration_ms=200.0)
|
|
assert ts[0] == 0
|
|
assert np.all(np.diff(ts) >= 1) # at least 1 ms apart
|
|
|
|
|
|
def test_build_timestamps_total_close_to_target() -> None:
|
|
log_dt = np.array([1.0] * 10)
|
|
ts = build_timestamps(log_dt, total_duration_ms=300.0)
|
|
# Last timestamp should be roughly total - one slot
|
|
assert abs(ts[-1] - 270) < 60 # tolerant of clipping
|
|
|
|
|
|
from ai_mouse._postprocess import sample_duration, truncnorm_sample
|
|
|
|
|
|
def test_truncnorm_sample_within_bounds() -> None:
|
|
rng = np.random.default_rng(0)
|
|
samples = [
|
|
truncnorm_sample(80.0, 30.0, 20.0, 300.0, rng) for _ in range(500)
|
|
]
|
|
arr = np.array(samples)
|
|
assert arr.min() >= 20.0
|
|
assert arr.max() <= 300.0
|
|
# Mean roughly close to mu
|
|
assert abs(arr.mean() - 80.0) < 5.0
|
|
|
|
|
|
def test_truncnorm_sample_far_outside_falls_back_to_clip() -> None:
|
|
rng = np.random.default_rng(0)
|
|
# mu far outside [low, high] — rejection will fail
|
|
v = truncnorm_sample(mu=1000.0, sigma=1.0, low=20.0, high=30.0, rng=rng)
|
|
assert 20.0 <= v <= 30.0
|
|
|
|
|
|
def test_sample_duration_uses_correct_bin() -> None:
|
|
dist_dict = {
|
|
"bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
|
|
"params": [
|
|
{"mu_log": 4.0, "sigma_log": 0.01}, # bin 0: dist < 50
|
|
{"mu_log": 5.0, "sigma_log": 0.01}, # bin 1: 50 <= dist < 100
|
|
{"mu_log": 6.0, "sigma_log": 0.01}, # bin 2: 100 <= dist < 200
|
|
] + [{"mu_log": 7.0, "sigma_log": 0.01}] * 5,
|
|
}
|
|
rng = np.random.default_rng(0)
|
|
v = sample_duration(dist_dict, 150.0, rng)
|
|
# exp(6) ~ 403, with tiny sigma we should land near there
|
|
assert 350 < v < 460
|
|
|
|
|
|
from ai_mouse._postprocess import soften_forward
|
|
|
|
|
|
def test_soften_forward_tolerates_small_backtrack() -> None:
|
|
# A 0.01 dip is within the 0.02 tolerance and must survive untouched.
|
|
f = np.array([0.0, 0.30, 0.29, 0.60, 1.0])
|
|
out = soften_forward(f)
|
|
assert np.isclose(out[2], 0.29)
|
|
|
|
|
|
def test_soften_forward_limits_large_backtrack() -> None:
|
|
# A 0.30 dip is noise; it gets pulled up to prev - tol.
|
|
f = np.array([0.0, 0.50, 0.20, 0.70, 1.0])
|
|
out = soften_forward(f)
|
|
assert np.isclose(out[2], 0.50 - 0.02)
|
|
|
|
|
|
def test_soften_forward_allows_moderate_overshoot() -> None:
|
|
# Overshoot past 1.0 is natural; small overshoot survives (compressed
|
|
# but strictly > 1.0).
|
|
f = np.array([0.0, 0.5, 0.9, 1.04, 1.0])
|
|
out = soften_forward(f)
|
|
assert out[3] > 1.0
|
|
|
|
|
|
def test_soften_forward_compresses_extreme_overshoot() -> None:
|
|
# tanh compression: no output value may exceed 1 + overshoot_span.
|
|
f = np.array([0.0, 0.5, 1.30, 1.50, 1.0])
|
|
out = soften_forward(f)
|
|
assert out.max() <= 1.0 + 0.08 + 1e-9
|
|
assert out[2] > 1.0 # still an overshoot, not clipped flat
|
|
|
|
|
|
def test_soften_forward_no_lower_clip() -> None:
|
|
# Small wind-up behind the start is allowed (warp_endpoints pins
|
|
# the first point later; interior may be slightly negative).
|
|
f = np.array([0.0, -0.01, 0.30, 0.70, 1.0])
|
|
out = soften_forward(f)
|
|
assert out[1] < 0.0
|
|
|
|
|
|
from ai_mouse._postprocess import damp_start
|
|
|
|
|
|
def test_damp_start_dampens_early_lateral() -> None:
|
|
lat = np.full(16, 1.0)
|
|
out = damp_start(lat, n=4)
|
|
assert out[1] < out[2] < out[3] < out[4] < 1.0 # monotone ramp
|
|
assert np.all(out[5:] == 1.0) # untouched past n
|
|
|
|
|
|
def test_damp_start_no_release_jump() -> None:
|
|
# The weight at i=n must be close to 1 (continuous release):
|
|
# smoothstep(4/5) = 0.896, vs the old linear 4/5 = 0.8.
|
|
lat = np.full(16, 1.0)
|
|
out = damp_start(lat, n=4)
|
|
assert out[4] > 0.85
|
|
|
|
|
|
def test_damp_start_short_input_safe() -> None:
|
|
lat = np.array([0.0, 0.5, 0.3])
|
|
out = damp_start(lat, n=4) # n capped to len//4 = 0 → no-op
|
|
assert np.array_equal(out, lat)
|
|
|
|
|
|
from ai_mouse._postprocess import warp_endpoints
|
|
|
|
|
|
def test_warp_endpoints_exact_pin() -> None:
|
|
f = np.linspace(0.05, 1.10, 32)
|
|
l = np.linspace(0.03, -0.07, 32)
|
|
fo, lo = warp_endpoints(f, l)
|
|
assert fo[0] == 0.0 and lo[0] == 0.0
|
|
assert fo[-1] == 1.0 and lo[-1] == 0.0
|
|
|
|
|
|
def test_warp_endpoints_identity_when_already_pinned() -> None:
|
|
f = np.linspace(0.0, 1.0, 32)
|
|
l = np.sin(np.linspace(0, np.pi, 32)) * 0.1
|
|
l[0] = l[-1] = 0.0
|
|
fo, lo = warp_endpoints(f.copy(), l.copy())
|
|
assert np.allclose(fo, f, atol=1e-12)
|
|
assert np.allclose(lo, l, atol=1e-12)
|
|
|
|
|
|
def test_warp_endpoints_preserves_smoothness() -> None:
|
|
# Correcting a smooth curve must not introduce sharp local bends:
|
|
# the warp adds a smoothstep-weighted offset, so the second
|
|
# difference (discrete curvature proxy) stays small.
|
|
f = np.linspace(0.02, 1.08, 32)
|
|
l = np.full(32, 0.05)
|
|
fo, lo = warp_endpoints(f, l)
|
|
assert np.abs(np.diff(lo, 2)).max() < 0.01
|
|
assert np.abs(np.diff(fo, 2)).max() < 0.01
|
|
|
|
|
|
def test_warp_endpoints_correction_local_to_each_end() -> None:
|
|
# A start-only residual should barely move the last quarter.
|
|
# (smoothstep weight at i=24/31 is ~0.13, so 0.08 residual leaves
|
|
# ~0.010 there — threshold 0.02 gives margin without losing meaning)
|
|
f = np.linspace(0.0, 1.0, 32) + 0.0
|
|
l = np.zeros(32)
|
|
f[0] = 0.08 # start residual only
|
|
fo, _ = warp_endpoints(f, l)
|
|
assert np.abs(fo[24:] - f[24:]).max() < 0.02
|