From 05be116cdecfe583f9ff0a761872e446f18b0ad7 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Sun, 10 May 2026 13:39:13 +0800 Subject: [PATCH] feat(eval): Markdown report builder with matplotlib plots Co-Authored-By: Claude Sonnet 4.6 --- ai_mouse/eval/report.py | 231 +++++++++++++++++++++++++++++++++++++ tests/test_eval_metrics.py | 42 +++++++ 2 files changed, 273 insertions(+) create mode 100644 ai_mouse/eval/report.py diff --git a/ai_mouse/eval/report.py b/ai_mouse/eval/report.py new file mode 100644 index 0000000..4bbafa7 --- /dev/null +++ b/ai_mouse/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 ai_mouse.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) diff --git a/tests/test_eval_metrics.py b/tests/test_eval_metrics.py index 7ffa2f4..2083447 100644 --- a/tests/test_eval_metrics.py +++ b/tests/test_eval_metrics.py @@ -111,3 +111,45 @@ class TestKlDivergence: 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) + + +class TestReportGeneration: + def test_generates_report_md(self, tmp_path): + """Smoke test: build_report writes an MD file with all expected sections.""" + from ai_mouse.eval.report import build_report + + # Synthetic generated traces (3 traces, 50 points each) + rng = np.random.default_rng(0) + gen_traces = [] + for _ in range(3): + xs = np.cumsum(rng.uniform(0, 5, 50)) + ys = np.cumsum(rng.uniform(-1, 1, 50)) + ts = np.cumsum(rng.uniform(5, 20, 50)) + gen_traces.append({"xs": xs, "ys": ys, "ts": ts}) + + # Synthetic reference + ref_traces = [] + for _ in range(5): + xs = np.cumsum(rng.uniform(0, 5, 50)) + ys = np.cumsum(rng.uniform(-1, 1, 50)) + ts = np.cumsum(rng.uniform(5, 20, 50)) + ref_traces.append({"xs": xs, "ys": ys, "ts": ts}) + + out_md = tmp_path / "report.md" + build_report( + generated_traces=gen_traces, + reference_traces=ref_traces, + output_md=out_md, + tag="smoke-test", + model_dir="/fake/model/dir", + ) + assert out_md.exists() + content = out_md.read_text(encoding="utf-8") + assert "# Eval Report" in content + assert "smoke-test" in content + assert "速度" in content or "speed" in content.lower() + assert "FFT" in content.upper() + # PNG plots should exist next to MD + plot_dir = tmp_path / "plots" + assert plot_dir.exists() + assert any(plot_dir.iterdir())