refactor: move eval/ and data_adapters/ to tools/

tools/eval/__init__.py

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

tools/eval/__main__.py

@@ -0,0 +1,127 @@
+"""CLI: `python -m ai_mouse.eval --model-dir ... --reference ... --output ...`
+
+Loads N synthetic start/end pairs, calls the generator, loads M reference
+traces from a Balabit-format jsonl, and writes a Markdown report.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]:
+    """Load up to n_samples reference traces from a JSONL file.
+
+    Returns list of {"xs","ys","ts"} 1-D ndarrays.
+    """
+    out: list[dict] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                rec = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            moves = [e for e in rec.get("events", []) if e.get("type") == "move"]
+            if len(moves) < 4:
+                continue
+            xs = np.array([e["x"] for e in moves], dtype=float)
+            ys = np.array([e["y"] for e in moves], dtype=float)
+            ts = np.array([e["t"] for e in moves], dtype=float)
+            out.append({"xs": xs, "ys": ys, "ts": ts})
+            if len(out) >= n_samples:
+                break
+    return out
+
+
+def _generate_n_samples(
+    model_dir: str, n_samples: int, seed: int = 0
+) -> list[dict]:
+    """Call the project's generator N times with random start/end pairs."""
+    from ai_mouse.generator import generate
+
+    rng = random.Random(seed)
+    out: list[dict] = []
+    for i in range(n_samples):
+        sx = rng.randint(50, 750)
+        sy = rng.randint(50, 550)
+        angle = rng.uniform(0, 2 * math.pi)
+        dist = rng.randint(100, 600)
+        ex = int(sx + dist * math.cos(angle))
+        ey = int(sy + dist * math.sin(angle))
+        ex = max(0, min(800, ex))
+        ey = max(0, min(600, ey))
+        try:
+            pts = generate(start=(sx, sy), end=(ex, ey), model_dir=model_dir)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("generate() failed at i=%d: %s", i, exc)
+            continue
+        # Drop click events (last 2)
+        moves = pts[:-2]
+        if len(moves) < 4:
+            continue
+        xs = np.array([p[0] for p in moves], dtype=float)
+        ys = np.array([p[1] for p in moves], dtype=float)
+        ts = np.array([p[2] for p in moves], dtype=float)
+        out.append({"xs": xs, "ys": ys, "ts": ts})
+    return out
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Generate eval report comparing model output to reference traces.")
+    parser.add_argument("--model-dir", required=True, help="Path to trained model dir (with flow_model.pt)")
+    parser.add_argument("--reference", type=Path, required=True, help="JSONL reference traces (Balabit holdout)")
+    parser.add_argument("--n-samples", type=int, default=200, help="Number of generated samples")
+    parser.add_argument("--n-reference", type=int, default=1000, help="Number of reference samples to load")
+    parser.add_argument("--output", type=Path, required=True, help="Output Markdown file")
+    parser.add_argument("--tag", default="eval", help="Tag string used in plot filenames")
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    if not Path(args.model_dir).exists():
+        logger.error("Model dir not found: %s", args.model_dir)
+        return 2
+    if not args.reference.exists():
+        logger.error("Reference jsonl not found: %s", args.reference)
+        return 2
+
+    logger.info("Loading reference from %s ...", args.reference)
+    ref_traces = _load_reference_jsonl(args.reference, args.n_reference)
+    logger.info("Loaded %d reference traces", len(ref_traces))
+
+    logger.info("Generating %d samples from %s ...", args.n_samples, args.model_dir)
+    gen_traces = _generate_n_samples(args.model_dir, args.n_samples, seed=args.seed)
+    logger.info("Generated %d valid traces", len(gen_traces))
+
+    if not gen_traces or not ref_traces:
+        logger.error("Empty trace sets — aborting")
+        return 1
+
+    from tools.eval.report import build_report
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    build_report(
+        generated_traces=gen_traces,
+        reference_traces=ref_traces,
+        output_md=args.output,
+        tag=args.tag,
+        model_dir=args.model_dir,
+    )
+    logger.info("Done. Report at %s", args.output)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

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

tools/eval/report.py

@@ -0,0 +1,231 @@
+"""Markdown report generation for eval results.
+
+Outputs a self-contained .md file with embedded PNG plots in a sibling
+'plots/' directory.
+"""
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+
+import matplotlib
+matplotlib.use("Agg")  # headless
+import matplotlib.pyplot as plt  # noqa: E402
+
+from tools.eval.metrics import (
+    compute_acceleration,
+    compute_jerk,
+    compute_speed,
+    compute_stats,
+    fft_spectrum,
+    kl_divergence_histograms,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _aggregate_kinematics(traces: list[dict]) -> dict[str, np.ndarray]:
+    """Concatenate per-trace speed/accel/jerk arrays from a list of traces.
+
+    Args:
+        traces: list of {"xs", "ys", "ts"} dicts (1-D ndarrays).
+
+    Returns:
+        dict with keys "speed", "accel", "jerk", "dt" — each a flat ndarray.
+    """
+    speeds, accels, jerks, dts = [], [], [], []
+    for tr in traces:
+        xs, ys, ts = tr["xs"], tr["ys"], tr["ts"]
+        if len(xs) < 4:
+            continue
+        v = compute_speed(xs, ys, ts)
+        a = compute_acceleration(v, ts)
+        j = compute_jerk(a, ts)
+        speeds.append(v)
+        accels.append(a)
+        jerks.append(j)
+        dts.append(np.diff(ts))
+    return {
+        "speed": np.concatenate(speeds) if speeds else np.array([]),
+        "accel": np.concatenate(accels) if accels else np.array([]),
+        "jerk": np.concatenate(jerks) if jerks else np.array([]),
+        "dt": np.concatenate(dts) if dts else np.array([]),
+    }
+
+
+def _plot_distribution(
+    gen: np.ndarray,
+    ref: np.ndarray,
+    title: str,
+    output: Path,
+    xlabel: str,
+    bins: int = 50,
+) -> None:
+    """Side-by-side histogram of gen vs ref."""
+    fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
+    if len(gen) > 0:
+        ax.hist(gen, bins=bins, alpha=0.5, label="生成", density=True)
+    if len(ref) > 0:
+        ax.hist(ref, bins=bins, alpha=0.5, label="参考 (Balabit)", density=True)
+    ax.set_title(title)
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel("密度")
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(output)
+    plt.close(fig)
+
+
+def _plot_fft_overlay(
+    gen_traces: list[dict],
+    ref_traces: list[dict],
+    output: Path,
+    sample_rate_hz: float = 100.0,
+) -> None:
+    """Average FFT magnitude over lateral component for gen vs ref."""
+    def _avg_spectrum(traces: list[dict]) -> tuple[np.ndarray, np.ndarray]:
+        all_freqs = None
+        all_mags = []
+        for tr in traces:
+            xs, ys = tr["xs"], tr["ys"]
+            if len(xs) < 8:
+                continue
+            sig = ys - np.linspace(ys[0], ys[-1], len(ys))
+            f, m = fft_spectrum(sig, sample_rate_hz)
+            if all_freqs is None:
+                all_freqs = f
+                all_mags.append(m)
+            elif len(m) == len(all_freqs):
+                all_mags.append(m)
+        if not all_mags:
+            return np.array([]), np.array([])
+        return all_freqs, np.mean(all_mags, axis=0)
+
+    fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
+    f_gen, m_gen = _avg_spectrum(gen_traces)
+    f_ref, m_ref = _avg_spectrum(ref_traces)
+    if len(f_gen) > 0:
+        ax.plot(f_gen, m_gen, label="生成", alpha=0.7)
+    if len(f_ref) > 0:
+        ax.plot(f_ref, m_ref, label="参考 (Balabit)", alpha=0.7)
+    ax.axvspan(4, 12, alpha=0.1, color="green", label="生理震颤区间 4–12 Hz")
+    ax.set_title("FFT 频谱（横向偏移信号）")
+    ax.set_xlabel("Hz")
+    ax.set_ylabel("|FFT|")
+    ax.set_xlim(0, sample_rate_hz / 2)
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(output)
+    plt.close(fig)
+
+
+def _plot_paths_overlay(traces: list[dict], output: Path, max_traces: int = 5) -> None:
+    """Plot up to N generated trajectories on the same axes."""
+    fig, ax = plt.subplots(figsize=(6, 5), dpi=100)
+    for i, tr in enumerate(traces[:max_traces]):
+        ax.plot(tr["xs"], tr["ys"], alpha=0.6, label=f"路径 {i+1}")
+    ax.invert_yaxis()  # screen coords
+    ax.set_title(f"前 {min(max_traces, len(traces))} 条生成轨迹")
+    ax.set_xlabel("x (px)")
+    ax.set_ylabel("y (px)")
+    ax.set_aspect("equal", adjustable="datalim")
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(output)
+    plt.close(fig)
+
+
+def build_report(
+    generated_traces: list[dict],
+    reference_traces: list[dict],
+    output_md: Path,
+    tag: str,
+    model_dir: str,
+    sample_rate_hz: float = 100.0,
+) -> None:
+    """Build a Markdown eval report with embedded plots.
+
+    Args:
+        generated_traces: list of {"xs","ys","ts"} from the generator under test.
+        reference_traces: list of {"xs","ys","ts"} from Balabit (ground truth).
+        output_md: destination .md path. plots/ created in same dir.
+        tag: short identifier (e.g. "baseline", "post-finetune").
+        model_dir: model directory path string (for provenance).
+        sample_rate_hz: nominal sample rate for FFT.
+    """
+    plot_dir = output_md.parent / "plots"
+    plot_dir.mkdir(parents=True, exist_ok=True)
+
+    gen_kin = _aggregate_kinematics(generated_traces)
+    ref_kin = _aggregate_kinematics(reference_traces)
+
+    # --- KL divergences ---
+    kl_speed = kl_divergence_histograms(gen_kin["speed"], ref_kin["speed"])
+    kl_accel = kl_divergence_histograms(gen_kin["accel"], ref_kin["accel"])
+    kl_jerk = kl_divergence_histograms(gen_kin["jerk"], ref_kin["jerk"])
+    kl_dt = kl_divergence_histograms(gen_kin["dt"], ref_kin["dt"])
+
+    # --- Stats ---
+    stats_gen = {k: compute_stats(v) for k, v in gen_kin.items()}
+    stats_ref = {k: compute_stats(v) for k, v in ref_kin.items()}
+
+    # --- Plots ---
+    _plot_distribution(gen_kin["speed"], ref_kin["speed"],
+                       "速度分布", plot_dir / f"{tag}-speed.png", "px/ms")
+    _plot_distribution(gen_kin["accel"], ref_kin["accel"],
+                       "加速度分布", plot_dir / f"{tag}-accel.png", "px/ms²")
+    _plot_distribution(gen_kin["jerk"], ref_kin["jerk"],
+                       "Jerk 分布", plot_dir / f"{tag}-jerk.png", "px/ms³")
+    _plot_distribution(gen_kin["dt"], ref_kin["dt"],
+                       "Δt 分布", plot_dir / f"{tag}-dt.png", "ms")
+    _plot_fft_overlay(generated_traces, reference_traces,
+                      plot_dir / f"{tag}-fft.png", sample_rate_hz)
+    _plot_paths_overlay(generated_traces, plot_dir / f"{tag}-paths.png")
+
+    # --- Markdown ---
+    now = datetime.now().strftime("%Y-%m-%d %H:%M")
+    lines = [
+        f"# Eval Report: {tag} ({now})",
+        "",
+        "## 模型信息",
+        f"- Checkpoint dir: `{model_dir}`",
+        f"- 生成样本数: {len(generated_traces)}",
+        f"- 参考样本数: {len(reference_traces)}",
+        "",
+        "## KL 散度（生成 vs 参考，越小越好）",
+        "| 指标 | KL |",
+        "|---|---|",
+        f"| 速度分布 | {kl_speed:.4f} |",
+        f"| 加速度分布 | {kl_accel:.4f} |",
+        f"| Jerk 分布 | {kl_jerk:.4f} |",
+        f"| Δt 分布 | {kl_dt:.4f} |",
+        "",
+        "## 摘要统计",
+        "| 指标 | 生成 mean | 参考 mean | 生成 CV | 参考 CV |",
+        "|---|---|---|---|---|",
+    ]
+    for key, label in [("speed", "速度"), ("accel", "加速度"), ("jerk", "jerk"), ("dt", "Δt")]:
+        lines.append(
+            f"| {label} | {stats_gen[key]['mean']:.4g} | {stats_ref[key]['mean']:.4g} | "
+            f"{stats_gen[key]['cv']:.3f} | {stats_ref[key]['cv']:.3f} |"
+        )
+    lines += [
+        "",
+        "## 直方图",
+        f"![速度](plots/{tag}-speed.png)",
+        f"![加速度](plots/{tag}-accel.png)",
+        f"![Jerk](plots/{tag}-jerk.png)",
+        f"![Δt](plots/{tag}-dt.png)",
+        "",
+        "## FFT 频谱（横向偏移）",
+        f"![FFT](plots/{tag}-fft.png)",
+        "",
+        "## 生成轨迹示例",
+        f"![轨迹](plots/{tag}-paths.png)",
+        "",
+    ]
+    output_md.write_text("\n".join(lines), encoding="utf-8")
+    logger.info("Report written to %s", output_md)