plot.py

#!/usr/bin/env python3
"""Compare the CMA-ES loss histories of two paper experiments."""

from __future__ import annotations

import argparse
import csv
import os
from pathlib import Path

os.environ.setdefault(
    "MPLCONFIGDIR",
    str(Path(os.environ.get("TMPDIR", "/tmp")) / "codesign_embedding_matplotlib"),
)

import matplotlib


PAPER_BLUE = "#2864aa"
PAPER_RED = "#c83c3c"


def load_cmaes_log(path: Path) -> tuple[list[int], list[float]]:
    """Load generations and best losses, joining sequential optimization phases."""
    generations: list[int] = []
    losses: list[float] = []
    generation_offset = 0
    previous_raw_generation: int | None = None

    with path.open(newline="", encoding="utf-8") as log_file:
        for line_number, row in enumerate(csv.reader(log_file), start=1):
            if not row or not any(field.strip() for field in row):
                continue

            try:
                raw_generation = int(row[0].strip())
                loss = float(row[2].strip())
            except (IndexError, ValueError):
                # Pagmo writes a header at the start of each appended run.
                if row[0].strip().lower() in {"gen", "generation"}:
                    continue
                raise ValueError(
                    f"Could not parse {path}:{line_number}; expected the CMA-ES "
                    "columns 'gen, fevals, best, dx, df, sigma'."
                ) from None

            if previous_raw_generation is not None and raw_generation <= previous_raw_generation:
                generation_offset = generations[-1]

            generations.append(raw_generation + generation_offset)
            losses.append(loss)
            previous_raw_generation = raw_generation

    if not generations:
        raise ValueError(f"No CMA-ES samples found in {path}")

    return generations, losses


def default_label(path: Path) -> str:
    return path.parent.name.replace("_", " ").replace("-", " ").title()


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Plot two cmaes.log files on one paper-style loss chart.",
    )
    parser.add_argument("log_a", type=Path, help="First cmaes.log file.")
    parser.add_argument("log_b", type=Path, help="Second cmaes.log file.")
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=Path("loss_comparison.pdf"),
        help="Output figure. The extension selects the format (default: loss_comparison.pdf).",
    )
    parser.add_argument(
        "--labels",
        nargs=2,
        metavar=("LABEL_A", "LABEL_B"),
        help="Legend labels (default: names of the logs' parent directories).",
    )
    parser.add_argument("--title", help="Optional plot title, for example 'Swimmer'.")
    parser.add_argument(
        "--colors",
        nargs=2,
        default=(PAPER_BLUE, PAPER_RED),
        metavar=("COLOR_A", "COLOR_B"),
        help="Matplotlib colors for the two curves (default: paper blue and red).",
    )
    parser.add_argument("--dpi", type=int, default=300, help="DPI for raster output.")
    parser.add_argument("--show", action="store_true", help="Also open the plot interactively.")
    return parser


def main() -> None:
    args = build_parser().parse_args()
    if not args.show:
        matplotlib.use("Agg")

    import matplotlib.pyplot as plt

    gen_a, loss_a = load_cmaes_log(args.log_a)
    gen_b, loss_b = load_cmaes_log(args.log_b)
    labels = args.labels or (default_label(args.log_a), default_label(args.log_b))

    # Match the compact square axes, colors, grid, and in-axis legend of the
    # paper's PGFPlots figures while remaining readable on a GitHub page.
    fig, ax = plt.subplots(figsize=(4.0, 4.0))
    ax.plot(gen_a, loss_a, color=args.colors[0], linewidth=1.8, label=labels[0])
    ax.plot(gen_b, loss_b, color=args.colors[1], linewidth=1.8, label=labels[1])
    ax.set_xlabel("Generation")
    ax.set_ylabel("Loss")
    ax.set_xlim(left=0)
    if args.title:
        ax.set_title(args.title)
    ax.grid(True, color="black", alpha=0.08, linewidth=0.6)
    ax.legend(
        loc="upper right",
        frameon=True,
        fancybox=True,
        framealpha=0.85,
        edgecolor="0.8",
        fontsize="small",
    )
    fig.tight_layout()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output, dpi=args.dpi, bbox_inches="tight")
    print(f"Saved loss comparison to: {args.output.resolve()}")

    if args.show:
        plt.show()
    else:
        plt.close(fig)


if __name__ == "__main__":
    main()