feat(eval): kinematics metrics, FFT spectrum, KL divergence

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ai_mouse/eval/__init__.py

@@ -0,0 +1 @@
+"""Evaluation module: kinematic metrics and Markdown report generation."""

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)))

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)