"""Parse FCOT-Separable training logs and visualize dual/grad/active metrics.""" import argparse import os import re from dataclasses import dataclass, field from datetime import datetime from typing import List, Dict, Any import matplotlib.pyplot as plt import numpy as np @dataclass class RunRecord: meta: Dict[str, Any] = field(default_factory=dict) iters: List[int] = field(default_factory=list) timestamps: List[str] = field(default_factory=list) dual: List[float] = field(default_factory=list) u: List[float] = field(default_factory=list) uc: List[float] = field(default_factory=list) grad: List[float] = field(default_factory=list) temp: List[float] = field(default_factory=list) active: List[float] = field(default_factory=list) react_steps: List[int] = field(default_factory=list) refresh_steps: List[int] = field(default_factory=list) ARCH_RE = re.compile( r"\[FCOT-SEP ARCH\] dim=(?P\d+), radius=(?P[\d\.]+), ny=(?P\d+)" ) ITER_RE = re.compile( r"\[(?P[^]]+)\] INFO:training: " r"\[Iter (?P\d+)\] dual=(?P[-0-9.e+]+) " r"u=(?P[-0-9.e+]+) uc=(?P[-0-9.e+]+) " r"grad=(?P[-0-9.e+]+).* temp=(?P[-0-9.e+]+) " r"active=(?P[-0-9.e+]+)%" ) REACT_RE = re.compile( r"\[FCOT-SEP\] Reactivated .* at step (?P\d+)" ) REFRESH_RE = re.compile( r"\[FCOT-SEP\] Full refresh updated .* at step (?P\d+)" ) TRAIN_PATH_RE = re.compile( r"tmp/FCOTSeparable_dim(?P\d+)_kernel-(?P.+?)_" ) def parse_training_log(path: str) -> List[RunRecord]: """ Scan `training.log` and extract iteration stats plus reactivation/refresh markers. Parameters ---------- path : str Path to the training log file generated by FCOT-Separable. Returns ------- List[RunRecord] Parsed run records sorted in chronological order. """ runs: List[RunRecord] = [] current: RunRecord | None = None with open(path, "r") as f: for line in f: line = line.strip() if not line: continue # Start of a new run: architecture line if "[FCOT-SEP ARCH]" in line: m = ARCH_RE.search(line) current = RunRecord() if m: current.meta["dim"] = int(m.group("dim")) current.meta["radius"] = float(m.group("radius")) current.meta["ny"] = int(m.group("ny")) runs.append(current) continue if current is None: # Skip lines before the first ARCH continue # Iteration line with dual/grad/etc. if "[Iter " in line: m = ITER_RE.search(line) if not m: continue current.timestamps.append(m.group("ts")) current.iters.append(int(m.group("iter"))) current.dual.append(float(m.group("dual"))) current.u.append(float(m.group("u"))) current.uc.append(float(m.group("uc"))) current.grad.append(float(m.group("grad"))) current.temp.append(float(m.group("temp"))) current.active.append(float(m.group("active"))) continue # Reactivation events if "[FCOT-SEP] Reactivated" in line: m = REACT_RE.search(line) if m: current.react_steps.append(int(m.group("step"))) continue # Full refresh events if "[FCOT-SEP] Full refresh updated" in line: m = REFRESH_RE.search(line) if m: current.refresh_steps.append(int(m.group("step"))) continue # Lines that mention checkpoint paths (TRAIN/CACHE/SAVE) → extract kernel if "[TRAIN] No checkpoint found at" in line or "[CACHE] Found saved" in line or "[SAVE] Writing checkpoint to" in line: m = TRAIN_PATH_RE.search(line) if m and current is not None: current.meta["kernel"] = m.group("kernel") continue return runs def plot_runs( runs: List[RunRecord], out_dir: str = "tmp/training_plots", ) -> None: os.makedirs(out_dir, exist_ok=True) for idx, run in enumerate(runs): fig, axes = plot_run(run, idx=idx, show=False) dim = run.meta.get("dim", "?") fname = os.path.join(out_dir, f"run_{idx}_dim{dim}.png") fig.savefig(fname, dpi=150) plt.close(fig) def plot_run(run: RunRecord, idx: int | None = None, show: bool = True): """ Plot a single RunRecord (dual, grad, active) and optionally show inline. Returns (fig, axes) so it can be used easily from notebooks. """ if not run.iters: raise ValueError("Run has no iteration data to plot.") iters = np.array(run.iters) dual = np.array(run.dual) grad = np.array(run.grad) active = np.array(run.active) # already in percent dim = run.meta.get("dim", "?") kernel = run.meta.get("kernel", "?") fig, axes = plt.subplots(3, 1, sharex=True, figsize=(12, 8)) # Duration estimate from first/last timestamp (if available) duration_str = "" if run.timestamps: try: t0 = datetime.strptime(run.timestamps[0], "%Y-%m-%d %H:%M:%S,%f") t1 = datetime.strptime(run.timestamps[-1], "%Y-%m-%d %H:%M:%S,%f") dt_min = (t1 - t0).total_seconds() / 60.0 duration_str = f", Δt≈{dt_min:.1f} min" except Exception: duration_str = "" # Dual title_suffix = f" (run {idx})" if idx is not None else "" axes[0].plot(iters, dual, label="dual") axes[0].set_ylabel("dual") axes[0].set_title(f"dim={dim}, kernel={kernel}{duration_str}{title_suffix}") # Gradient norm axes[1].plot(iters, grad, label="grad_norm", color="tab:orange") axes[1].set_ylabel("grad norm") # Active fraction axes[2].plot(iters, active, label="active %", color="tab:green") axes[2].set_ylabel("active (%)") axes[2].set_xlabel("iteration") # Reactivation / refresh markers for ax in axes: for s in run.react_steps: ax.axvline(s, color="red", linewidth=0.6, alpha=0.6) for s in run.refresh_steps: ax.axvline(s, color="blue", linewidth=0.6, alpha=0.6) axes[0].legend(loc="best") axes[1].legend(loc="best") axes[2].legend(loc="best") fig.tight_layout() if show: plt.show() return fig, axes def main(): """Entry point to parse `training.log` and dump summary plots to disk.""" parser = argparse.ArgumentParser( description="Parse training.log and plot dual/grad/active with reactivation/refresh markers." ) parser.add_argument( "--log", type=str, default="training.log", help="Path to training log (default: training.log)", ) parser.add_argument( "--out", type=str, default="tmp/training_plots", help="Output directory for plots (default: tmp/training_plots)", ) args = parser.parse_args() runs = parse_training_log(args.log) if not runs: print(f"No runs found in {args.log}") return print(f"Parsed {len(runs)} runs from {args.log}") plot_runs(runs, out_dir=args.out) print(f"Saved plots to {args.out}") if __name__ == "__main__": main()