158 lines
4.8 KiB
Python
158 lines
4.8 KiB
Python
"""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)))
|