Files
ai_mouse/tools/eval/report.py

232 lines
7.8 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""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="生理震颤区间 412 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)