scripts/make_visibility_paper_figure.py

import argparse
import json
from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np


FEATURE_LABELS = {
    "angle": "Incidence angle (degrees)",
    "distance": "View-to-point distance (meters)",
    "contrast": "Local contrast",
    "blur": "Blur response",
    "snr": "Signal-to-noise ratio",
    "saturation": "Saturation",
}

FEATURE_ORDER = ["angle", "distance", "contrast", "blur", "snr", "saturation"]
FEATURE_GROUPS = {
    "geometry": ["angle", "distance"],
    "image": ["contrast", "blur"],
}
FEATURE_COLORS = {
    "angle": "#1f77b4",
    "distance": "#d62728",
    "contrast": "#2ca02c",
    "blur": "#9467bd",
    "snr": "#8c564b",
    "saturation": "#ff7f0e",
}


def setup_style():
    mpl.rcParams.update(
        {
            "figure.dpi": 180,
            "savefig.dpi": 300,
            "font.family": "sans-serif",
            "font.sans-serif": ["Helvetica", "Arial", "DejaVu Sans"],
            "mathtext.fontset": "dejavusans",
            "axes.spines.top": False,
            "axes.spines.right": False,
            "axes.linewidth": 0.8,
            "axes.labelsize": 8,
            "axes.titlesize": 9,
            "xtick.labelsize": 7,
            "ytick.labelsize": 7,
            "legend.fontsize": 7,
            "grid.alpha": 0.12,
            "grid.linewidth": 0.4,
            "lines.linewidth": 1.2,
        }
    )


def _load_summary(path):
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def _load_curve_json(path):
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def _feature_curve(feature_payload):
    curve = feature_payload["curve"]
    x = np.array([row["x_center"] for row in curve], dtype=np.float64)
    y = np.array([row["weight_mean"] for row in curve], dtype=np.float64)
    if curve and "weight_q25" in curve[0] and "weight_q75" in curve[0]:
        low = np.array([row["weight_q25"] for row in curve], dtype=np.float64)
        high = np.array([row["weight_q75"] for row in curve], dtype=np.float64)
    else:
        err = np.array([row["weight_stderr"] for row in curve], dtype=np.float64)
        low = y - err
        high = y + err
    return x, y, low, high


def _global_y_limits(summary):
    y_values = []
    for payload in summary["features"].values():
        for row in payload["curve"]:
            y_values.append(float(row["weight_mean"]))
    if not y_values:
        return 0.0, 1.0
    ymin = min(y_values)
    ymax = max(y_values)
    pad = max((ymax - ymin) * 0.08, 0.01)
    return max(0.0, ymin - pad), min(1.0, ymax + pad)


def _resolve_y_limits(summary, ymin=None, ymax=None):
    auto_ymin, auto_ymax = _global_y_limits(summary)
    return (
        auto_ymin if ymin is None else float(ymin),
        auto_ymax if ymax is None else float(ymax),
    )


def _plot_panel(ax, feature_name, feature_payload, show_ylabel=True):
    x, y, low, high = _feature_curve(feature_payload)
    color = FEATURE_COLORS[feature_name]
    xlabel = FEATURE_LABELS.get(feature_name, feature_name.title())

    line = ax.plot(
        x,
        y,
        color=color,
        marker="o",
        markersize=1.8,
        markeredgewidth=0.0,
        label="Mean",
    )[0]
    band = ax.fill_between(
        x,
        low,
        high,
        color=color,
        alpha=0.10,
        linewidth=0,
        label="IQR (25--75%)",
    )
    ax.set_xlabel(xlabel)
    if show_ylabel:
        ax.set_ylabel("Mean learned weight")
    ax.grid(True, axis="y")
    ax.margins(x=0.01)
    ax.legend(handles=[line, band], loc="best", frameon=False, handlelength=2.0)


def _save_individual_figure(feature_name, feature_payload, output_dir, y_limits):
    fig, ax = plt.subplots(figsize=(3.35, 2.45), constrained_layout=True)
    _plot_panel(ax, feature_name, feature_payload)
    ax.set_ylim(*y_limits)
    pdf_path = output_dir / f"{feature_name}.pdf"
    png_path = output_dir / f"{feature_name}.png"
    fig.savefig(pdf_path, bbox_inches="tight")
    fig.savefig(png_path, bbox_inches="tight")
    plt.close(fig)
    return pdf_path, png_path


def _save_group_figure(summary, feature_names, output_prefix, y_limits):
    fig, axes = plt.subplots(1, len(feature_names), figsize=(7.0, 2.6), constrained_layout=True)
    if len(feature_names) == 1:
        axes = [axes]
    plotted = 0
    for idx, (ax, feature_name) in enumerate(zip(axes, feature_names)):
        if feature_name not in summary["features"]:
            ax.axis("off")
            continue
        feature_payload = summary["features"][feature_name]
        _plot_panel(ax, feature_name, feature_payload, show_ylabel=(idx == 0))
        ax.set_ylim(*y_limits)
        plotted += 1
    pdf_path = output_prefix.with_suffix(".pdf")
    png_path = output_prefix.with_suffix(".png")
    fig.savefig(pdf_path, bbox_inches="tight")
    fig.savefig(png_path, bbox_inches="tight")
    plt.close(fig)
    return pdf_path, png_path, plotted


def _save_geometry_paper_figure(summary, output_prefix, y_limits, angle_payload=None):
    feature_names = ["angle", "distance"]
    fig, axes = plt.subplots(1, 2, figsize=(6.8, 2.55), constrained_layout=True)

    for idx, (ax, feature_name) in enumerate(zip(axes, feature_names)):
        if feature_name == "angle" and angle_payload is not None:
            feature_payload = angle_payload
        else:
            feature_payload = summary["features"][feature_name]
        _plot_panel(ax, feature_name, feature_payload, show_ylabel=(idx == 0))
        ax.set_ylim(*y_limits)

    pdf_path = output_prefix.with_suffix(".pdf")
    png_path = output_prefix.with_suffix(".png")
    fig.savefig(pdf_path, bbox_inches="tight")
    fig.savefig(png_path, bbox_inches="tight")
    plt.close(fig)
    return pdf_path, png_path


def _load_angle_override(summary_path, class_id):
    if class_id is None:
        return None
    curve_path = summary_path.parent / "per_class" / f"class_{class_id}" / "angle_curve.json"
    if not curve_path.exists():
        raise FileNotFoundError(f"Missing per-class angle curve: {curve_path}")
    payload = _load_curve_json(curve_path)
    return {"curve": payload["curve"]}


def build_figure(summary, summary_path, output_prefix, ymin=None, ymax=None, angle_class=None):
    setup_style()
    fig, axes = plt.subplots(2, 3, figsize=(10.6, 5.9), constrained_layout=True)
    axes = axes.ravel()
    individual_dir = output_prefix.parent / f"{output_prefix.stem}_individual"
    individual_dir.mkdir(parents=True, exist_ok=True)
    y_limits = _resolve_y_limits(summary, ymin=ymin, ymax=ymax)
    angle_override = _load_angle_override(summary_path, angle_class)

    plotted = 0
    for ax, feature_name in zip(axes, FEATURE_ORDER):
        if feature_name not in summary["features"]:
            ax.axis("off")
            continue
        if feature_name == "angle" and angle_override is not None:
            feature_payload = angle_override
        else:
            feature_payload = summary["features"][feature_name]
        _plot_panel(ax, feature_name, feature_payload)
        ax.set_ylim(*y_limits)
        _save_individual_figure(feature_name, feature_payload, individual_dir, y_limits)
        plotted += 1

    pdf_path = output_prefix.with_suffix(".pdf")
    png_path = output_prefix.with_suffix(".png")
    fig.savefig(pdf_path, bbox_inches="tight")
    fig.savefig(png_path, bbox_inches="tight")
    plt.close(fig)

    group_outputs = {}
    for group_name, feature_names in FEATURE_GROUPS.items():
        group_prefix = output_prefix.parent / f"{output_prefix.stem}_{group_name}"
        group_pdf, group_png, _ = _save_group_figure(summary, feature_names, group_prefix, y_limits)
        group_outputs[group_name] = {"pdf": str(group_pdf), "png": str(group_png)}

    geometry_clean_prefix = output_prefix.parent / f"{output_prefix.stem}_geometry_clean"
    geometry_clean_pdf, geometry_clean_png = _save_geometry_paper_figure(
        summary,
        geometry_clean_prefix,
        y_limits,
        angle_payload=angle_override,
    )
    group_outputs["geometry_clean"] = {"pdf": str(geometry_clean_pdf), "png": str(geometry_clean_png)}

    return pdf_path, png_path, plotted


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary", type=Path, required=True)
    parser.add_argument("--output_prefix", type=Path, required=True)
    parser.add_argument("--ymin", type=float, default=None)
    parser.add_argument("--ymax", type=float, default=None)
    parser.add_argument("--angle_class", type=int, default=None)
    args = parser.parse_args()

    summary = _load_summary(args.summary)
    args.output_prefix.parent.mkdir(parents=True, exist_ok=True)
    pdf_path, png_path, plotted = build_figure(
        summary,
        args.summary,
        args.output_prefix,
        ymin=args.ymin,
        ymax=args.ymax,
        angle_class=args.angle_class,
    )
    print(f"Saved {plotted} panels to {pdf_path} and {png_path}")


if __name__ == "__main__":
    main()