From dc38f031b8ce09ff1ea4459a8c13abe057156499 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Sun, 10 May 2026 13:33:19 +0800 Subject: [PATCH] feat(eval): kinematics metrics, FFT spectrum, KL divergence Co-Authored-By: Claude Sonnet 4.6 --- ai_mouse/eval/__init__.py | 1 + ai_mouse/eval/metrics.py | 157 +++++++++++++++++++++++++++++++++++++ tests/test_eval_metrics.py | 113 ++++++++++++++++++++++++++ 3 files changed, 271 insertions(+) create mode 100644 ai_mouse/eval/__init__.py create mode 100644 ai_mouse/eval/metrics.py create mode 100644 tests/test_eval_metrics.py diff --git a/ai_mouse/eval/__init__.py b/ai_mouse/eval/__init__.py new file mode 100644 index 0000000..af79cb6 --- /dev/null +++ b/ai_mouse/eval/__init__.py @@ -0,0 +1 @@ +"""Evaluation module: kinematic metrics and Markdown report generation.""" diff --git a/ai_mouse/eval/metrics.py b/ai_mouse/eval/metrics.py new file mode 100644 index 0000000..ddc4b2a --- /dev/null +++ b/ai_mouse/eval/metrics.py @@ -0,0 +1,157 @@ +"""Kinematic metrics for mouse trajectory evaluation. + +All inputs are 1-D NumPy arrays. Time is in milliseconds, position in pixels. +Velocities are px/ms, accelerations px/ms², jerks px/ms³. +""" +from __future__ import annotations + +import numpy as np + + +def compute_speed( + xs: np.ndarray, ys: np.ndarray, ts: np.ndarray, eps: float = 1e-6 +) -> np.ndarray: + """Compute scalar speed at each step. + + Args: + xs: (N,) x coordinates. + ys: (N,) y coordinates. + ts: (N,) timestamps in ms. + eps: minimum dt (ms) to avoid div-by-zero. + + Returns: + (N-1,) array of speeds (px/ms). + """ + dx = np.diff(xs) + dy = np.diff(ys) + dt = np.maximum(np.diff(ts), eps) + return np.hypot(dx, dy) / dt + + +def compute_acceleration(speeds: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray: + """Compute scalar acceleration from speeds. + + Args: + speeds: (M,) speeds (px/ms). Typically M = N-1 from compute_speed. + ts: (M,) or (M+1,) timestamps in ms. + If len(ts) == len(speeds): timestamps are treated as the time + points associated with each speed value directly. + If len(ts) == len(speeds)+1: timestamps are the original position + timestamps; midpoints are computed for speed intervals. + eps: minimum dt (ms) to avoid div-by-zero. + + Returns: + (M-1,) array of accelerations (px/ms²). + """ + if len(speeds) < 2: + return np.array([], dtype=float) + if len(ts) == len(speeds): + # ts[i] is already the time associated with speed[i] + dt = np.maximum(np.diff(ts), eps) + else: + # ts has length M+1; speed[i] is between ts[i] and ts[i+1] + midpoints = (ts[:-1] + ts[1:]) / 2.0 + dt = np.maximum(np.diff(midpoints), eps) + return np.diff(speeds) / dt + + +def compute_jerk(accels: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray: + """Compute jerk from accelerations. + + Args: + accels: (K,) accelerations. + ts: (K+2,) timestamps that produced those accelerations. + Used to derive midpoint-of-midpoint dts. + eps: minimum dt to avoid div-by-zero. + + Returns: + (K-1,) array of jerks (px/ms³). + """ + if len(accels) < 2: + return np.array([], dtype=float) + # Approximate dt for jerks as average dt of original ts (good enough for stats) + dt_avg = np.maximum(np.diff(ts).mean(), eps) + return np.diff(accels) / dt_avg + + +def compute_stats(x: np.ndarray) -> dict[str, float]: + """Summary statistics for a 1-D distribution. + + Returns: + dict with keys: mean, std, cv (coef of variation), p25, p50, p75, p95. + """ + if len(x) == 0: + return {k: 0.0 for k in ("mean", "std", "cv", "p25", "p50", "p75", "p95")} + x = np.asarray(x, dtype=float) + mean = float(x.mean()) + std = float(x.std(ddof=1)) if len(x) > 1 else 0.0 + cv = std / mean if mean != 0 else 0.0 + return { + "mean": mean, + "std": std, + "cv": cv, + "p25": float(np.percentile(x, 25)), + "p50": float(np.percentile(x, 50)), + "p75": float(np.percentile(x, 75)), + "p95": float(np.percentile(x, 95)), + } + + +def fft_spectrum( + signal: np.ndarray, sample_rate_hz: float +) -> tuple[np.ndarray, np.ndarray]: + """Compute one-sided FFT magnitude spectrum. + + Args: + signal: 1-D real-valued signal. + sample_rate_hz: Sampling rate in Hz. + + Returns: + (freqs, magnitudes) — positive frequencies only. + Magnitudes are absolute values of complex FFT coefficients. + """ + n = len(signal) + if n == 0: + return np.array([]), np.array([]) + # Zero-mean to remove DC component which dominates the spectrum + s = signal - signal.mean() + fft = np.fft.rfft(s) + freqs = np.fft.rfftfreq(n, d=1.0 / sample_rate_hz) + return freqs, np.abs(fft) + + +def kl_divergence_histograms( + x: np.ndarray, + y: np.ndarray, + bins: int = 50, + eps: float = 1e-10, +) -> float: + """KL divergence KL(P_x || P_y) estimated via shared-bin histograms. + + Both arrays are histogrammed over their joint range. Empty bins get + `eps` mass to avoid log(0) — keeps result finite even for disjoint + supports. + + Args: + x: samples from distribution P. + y: samples from distribution Q (the "reference"). + bins: number of histogram bins. + eps: smoothing constant for empty bins. + + Returns: + scalar KL divergence (nats). Always finite, ≥ 0. + """ + if len(x) == 0 or len(y) == 0: + return 0.0 + lo = float(min(x.min(), y.min())) + hi = float(max(x.max(), y.max())) + if hi <= lo: + return 0.0 + edges = np.linspace(lo, hi, bins + 1) + px, _ = np.histogram(x, bins=edges, density=False) + qy, _ = np.histogram(y, bins=edges, density=False) + px = px.astype(float) + eps + qy = qy.astype(float) + eps + px /= px.sum() + qy /= qy.sum() + return float(np.sum(px * np.log(px / qy))) diff --git a/tests/test_eval_metrics.py b/tests/test_eval_metrics.py new file mode 100644 index 0000000..7ffa2f4 --- /dev/null +++ b/tests/test_eval_metrics.py @@ -0,0 +1,113 @@ +"""Tests for the eval metrics module.""" +from __future__ import annotations + +import numpy as np +import pytest + + +class TestKinematics: + def test_compute_speed_constant_velocity(self): + """Constant-velocity trajectory has constant speed.""" + from ai_mouse.eval.metrics import compute_speed + # 10 points, moving 10 px in 100 ms each step → speed = 0.1 px/ms + xs = np.arange(0, 100, 10, dtype=float) + ys = np.zeros(10, dtype=float) + ts = np.arange(0, 1000, 100, dtype=float) + v = compute_speed(xs, ys, ts) + # All speeds should be ≈ 0.1 px/ms + assert v.shape == (9,) # n-1 differences + np.testing.assert_allclose(v, 0.1, rtol=1e-4) + + def test_compute_speed_handles_zero_dt(self): + """Adjacent points with same timestamp must not produce NaN/inf.""" + from ai_mouse.eval.metrics import compute_speed + xs = np.array([0.0, 10.0, 20.0]) + ys = np.array([0.0, 0.0, 0.0]) + ts = np.array([0.0, 0.0, 100.0]) # zero dt between [0] and [1] + v = compute_speed(xs, ys, ts) + assert np.isfinite(v).all() + + def test_compute_acceleration(self): + """Linearly increasing speed → constant acceleration.""" + from ai_mouse.eval.metrics import compute_acceleration + # speeds: 0.1, 0.2, 0.3, 0.4 over dt = 100 ms each → a = 0.001 px/ms² + speeds = np.array([0.1, 0.2, 0.3, 0.4]) + ts = np.array([100.0, 200.0, 300.0, 400.0]) + a = compute_acceleration(speeds, ts) + np.testing.assert_allclose(a, 0.001, rtol=1e-4) + + def test_compute_jerk(self): + from ai_mouse.eval.metrics import compute_jerk + # accelerations: 0.001, 0.002, 0.003 over dt = 100 ms → j = 0.00001 + accels = np.array([0.001, 0.002, 0.003]) + ts = np.array([200.0, 300.0, 400.0]) + j = compute_jerk(accels, ts) + np.testing.assert_allclose(j, 1e-5, rtol=1e-4) + + +class TestStatsSummary: + def test_compute_stats_returns_expected_keys(self): + from ai_mouse.eval.metrics import compute_stats + x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + s = compute_stats(x) + assert "mean" in s + assert "std" in s + assert "cv" in s + assert "p25" in s + assert "p50" in s + assert "p75" in s + assert "p95" in s + + def test_cv_for_constant_is_zero(self): + from ai_mouse.eval.metrics import compute_stats + x = np.full(10, 3.0) + s = compute_stats(x) + assert s["cv"] == 0.0 + + +class TestFftSpectrum: + def test_finds_dominant_frequency(self): + """A pure 8 Hz signal should have its peak near 8 Hz.""" + from ai_mouse.eval.metrics import fft_spectrum + # Sample at 100 Hz for 1 second + sample_rate_hz = 100.0 + ts_ms = np.arange(0, 1000, 1000 / sample_rate_hz) + signal = np.sin(2 * np.pi * 8 * ts_ms / 1000) # 8 Hz sine + freqs, mags = fft_spectrum(signal, sample_rate_hz) + peak_freq = freqs[np.argmax(mags)] + assert abs(peak_freq - 8.0) < 1.0 # within 1 Hz + + def test_returns_only_positive_frequencies(self): + from ai_mouse.eval.metrics import fft_spectrum + signal = np.random.randn(64) + freqs, mags = fft_spectrum(signal, 50.0) + assert (freqs >= 0).all() + assert len(freqs) == len(mags) + + +class TestKlDivergence: + def test_identical_distributions_zero_kl(self): + """KL(p, p) ≈ 0.""" + from ai_mouse.eval.metrics import kl_divergence_histograms + rng = np.random.default_rng(42) + x = rng.normal(0, 1, 5000) + y = rng.normal(0, 1, 5000) + kl = kl_divergence_histograms(x, y, bins=50) + assert kl < 0.05 + + def test_different_distributions_positive_kl(self): + """Different means → positive KL.""" + from ai_mouse.eval.metrics import kl_divergence_histograms + rng = np.random.default_rng(42) + x = rng.normal(0, 1, 5000) + y = rng.normal(3, 1, 5000) + kl = kl_divergence_histograms(x, y, bins=50) + assert kl > 0.5 + + def test_handles_disjoint_supports(self): + """No NaN even when histograms have non-overlapping bins.""" + from ai_mouse.eval.metrics import kl_divergence_histograms + x = np.array([1.0, 1.1, 1.2, 1.3, 1.4]) + y = np.array([10.0, 10.1, 10.2, 10.3, 10.4]) + kl = kl_divergence_histograms(x, y, bins=10) + assert np.isfinite(kl)