MeanField/experiments/plot_coupled_convergence.py

#!/usr/bin/env python3
"""Plot completed, artifact-passed coupled-convergence audit summaries.

Example (run rendering serially with numerical campaigns):
  python plot_coupled_convergence.py analysis.json --output NEW_FIGURE_DIRECTORY
Use --pair n3 or --pair n2 to request exactly one complete pair. Default: both.
The script reads summary JSON only; it does not rerun the underlying FEM/artifact
audit. Matplotlib is imported only when figures are actually rendered.
"""
from __future__ import annotations

import argparse
from dataclasses import dataclass
import hashlib
import json
import math
from pathlib import Path
import sys


PAIRS = {
    "n3": ("n3_J0_control", "n3_J0_tight"),
    "n2": ("n2_restart_raw", "n2_restart_gravity_path"),
}
LABELS = {
    "n3_J0_control": "Control linear tolerance",
    "n3_J0_tight": "Tighter linear tolerance",
    "n2_restart_raw": "Raw trial path",
    "n2_restart_gravity_path": "Gravity-completed trial path",
}


def require(condition, message):
    if not condition:
        raise ValueError(message)


def finite(value, name):
    number = float(value)
    require(math.isfinite(number), f"Nonfinite {name}")
    return number


def near(a, b, name):
    require(math.isclose(a, b, rel_tol=3e-11, abs_tol=0.), f"Inconsistent {name}: {a} versus {b}")


def digest_bytes(data):
    return hashlib.sha256(data).hexdigest()


@dataclass
class Curve:
    identifier: str
    campaign: str
    audit_path: str
    producer_commit: str
    mesh_sha256: str
    metadata: dict
    input_state_hashes: dict
    time: list[float]
    residual: list[float]
    accepted: list[bool]
    threshold: float
    outcome: str
    failure: str | None


def curve(case, campaign, audit_path):
    name = case["id"]
    require(case.get("complete") is True and case.get("artifact_passed") is True,
            f"Case {name} is incomplete or did not pass the artifact audit")
    require(case.get("kind") == "newton" and case.get("outcome") in ("converged", "not_converged"),
            f"Case {name} has no completed Newton trajectory")
    initial = finite(case["initial_residual"], "initial residual")
    require(initial > 0, "Log plots require strictly positive residuals; exact zeros are not silently floored")
    atol, rtol = finite(case["atol"], "atol"), finite(case["rtol"], "rtol")
    require(atol >= 0 and rtol >= 0, "Nonlinear tolerances must be nonnegative")
    threshold = finite(case["actual_threshold"], "actual stopping threshold")
    near(threshold, max(atol, rtol * initial), "actual stopping threshold")
    require(threshold > 0, "A zero stopping threshold cannot be drawn on an unmodified log scale")
    history = case["history"]
    require(isinstance(history, list) and len(history) == int(case["attempted_updates"]), "Missing Newton history rows")
    times, residuals, accepted, durations = [0.], [initial], [], []
    for index, row in enumerate(history):
        require(int(row["iteration"]) == index, "Missing, repeated or unordered Newton iteration")
        require(str(row["accepted"]) in ("0", "1"), "Invalid acceptance flag")
        seconds = finite(row["seconds"], "Newton duration")
        residual = finite(row["F"], "Newton residual")
        require(seconds >= 0 and residual > 0, "Negative time or nonpositive log residual")
        is_accepted = str(row["accepted"]) == "1"
        if not is_accepted:
            near(residual, residuals[-1], "restored residual after failed Newton attempt")
        durations.append(seconds)
        times.append(math.fsum(durations))
        residuals.append(residual)
        accepted.append(is_accepted)
    near(times[-1], finite(case["newton_seconds"], "reported Newton time"), "summed Newton time")
    near(residuals[-1], finite(case["final_residual"], "final residual"), "final residual")
    require(sum(accepted) == int(case["accepted_updates"]), "Accepted-update count differs from history")
    require(case["outcome"] != "converged" or residuals[-1] <= threshold,
            "Convergence outcome exceeds the actual stopping threshold")
    return Curve(name, campaign["path"], audit_path, campaign["producer_commit"], campaign["mesh_sha256"],
                 case["metadata"], case.get("run", {}).get("input_state_hashes", {}),
                 times, residuals, accepted, threshold, case["outcome"], case.get("failure"))


def load_pairs(paths, selected="both"):
    require(selected in ("both", "n3", "n2"), "Unknown requested curve pair")
    requested = list(PAIRS) if selected == "both" else [selected]
    wanted = {name for pair in requested for name in PAIRS[pair]}
    curves, sources = {}, []
    for path in paths:
        path = Path(path).resolve()
        raw = path.read_bytes()
        audit = json.loads(raw)
        require(audit.get("complete") is True and audit.get("artifact_passed") is True,
                f"Only completed, artifact-passed audit summaries may be plotted: {path}")
        sources.append({"path": str(path), "sha256": digest_bytes(raw),
                        "auditor_sha256": audit["auditor_sha256"],
                        "benchmark_checker_sha256": audit.get("benchmark_checker_sha256"),
                        "audit_scope": audit.get("scope"), "audit_limitations": audit.get("limitations", [])})
        for campaign in audit["campaigns"]:
            require(campaign.get("complete") is True and campaign.get("artifact_passed") is True,
                    "An audit contains an incomplete or failed campaign")
            for case in campaign["cases"]:
                if case.get("id") in wanted:
                    require(case["id"] not in curves, f"Ambiguous duplicate case {case['id']} across audits")
                    curves[case["id"]] = curve(case, campaign, str(path))
    require(set(curves) == wanted, f"Missing requested curves: {sorted(wanted - set(curves))}")
    pairs = {}
    for name in requested:
        left, right = (curves[key] for key in PAIRS[name])
        require(left.mesh_sha256 == right.mesh_sha256, "A paired comparison uses different mesh hashes")
        for key in ("n", "J", "K", "rho_c", "surface_prescription", "surface_degree", "surface_maximum_order",
                    "density_order", "enthalpy_order", "gravity_potential_order", "gravity_flux_order", "displacement_order"):
            require(left.metadata.get(key) == right.metadata.get(key), f"Paired physical/discrete setting differs: {key}")
        require(float(left.metadata["n"]) == (3. if name == "n3" else 2.) and float(left.metadata["J"]) == 0.,
                "A named control pair has unexpected physical parameters")
        if name == "n2":
            require(bool(left.input_state_hashes) and left.input_state_hashes == right.input_state_hashes,
                    "Restart pair does not declare the same hash-bound input state files")
            near(left.residual[0], right.residual[0], "shared restart initial norm")
        pairs[name] = (left, right)
    return pairs, sources


def residual_reduction(candidate):
    return [100. * (1. - value / candidate.residual[0]) for value in candidate.residual]


def render_pair(name, curves):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    local = None
    if name == "n2":
        figure, (axis, local) = plt.subplots(2, 1, sharex=True, figsize=(8.5, 8.3),
                                           gridspec_kw={"height_ratios": [1.35, 1.]}, layout="constrained")
    else:
        figure, axis = plt.subplots(figsize=(8.5, 5.7), layout="constrained")
    colors = ("#1769a8", "#c35b16")
    markers = ("o", "s")
    for candidate, color, marker in zip(curves, colors, markers):
        label = LABELS[candidate.identifier]
        tolerance = candidate.metadata.get("linear_tol")
        if tolerance is not None:
            label += f" (linear rtol {float(tolerance):g})"
        label += "; " + ("converged" if candidate.outcome == "converged" else "not converged")
        axis.plot(candidate.time, candidate.residual, color=color, marker=marker, markersize=4,
                  linewidth=1.6, label=label)
        failed = [i + 1 for i, accepted in enumerate(candidate.accepted) if not accepted]
        if failed:
            axis.scatter([candidate.time[i] for i in failed], [candidate.residual[i] for i in failed],
                         marker="x", s=65, linewidths=1.6, color=color, zorder=4)
        if local is not None:
            reduction = residual_reduction(candidate)
            local.plot(candidate.time, reduction, color=color, marker=marker, markersize=4, linewidth=1.6)
            if failed:
                local.scatter([candidate.time[i] for i in failed], [reduction[i] for i in failed],
                              marker="x", s=65, linewidths=1.6, color=color, zorder=4)
    if curves[0].threshold == curves[1].threshold:
        axis.axhline(curves[0].threshold, color="#555555", linestyle=":", linewidth=1.4,
                     label=f"Actual stopping threshold (both): {curves[0].threshold:.4g}")
    else:
        for candidate, color in zip(curves, colors):
            axis.axhline(candidate.threshold, color=color, linestyle=":", linewidth=1.3,
                         label=f"{LABELS[candidate.identifier]} threshold: {candidate.threshold:.4g}")
    axis.set_yscale("log")
    axis.set_xlim(left=0., right=max(1., max(c.time[-1] for c in curves) * 1.04))
    if local is None:
        axis.set_xlabel("Cumulative reported Newton time (s)")
    else:
        # The full panel retains the stopping threshold. A separate zero-based
        # local quantity reveals percent-level progress without compressing or
        # hiding the many orders of magnitude still required for convergence.
        all_reductions = [value for candidate in curves for value in residual_reduction(candidate)]
        low, high = min(0., min(all_reductions)), max(0., max(all_reductions))
        padding = .08 * max(high - low, .1)
        local.set_ylim(low - padding if low < 0 else 0., high + padding)
        local.set_ylabel("Residual reduction from\nshared restart (%)")
        local.set_xlabel("Cumulative reported Newton time (s)")
        local.set_title("Local progress from the same trajectories", fontsize=10)
        local.grid(True, alpha=.22)
        local.text(.02, .04, "Local reduction scale; actual stopping thresholds remain in the full-scale panel above.",
                   fontsize=8, transform=local.transAxes,
                   bbox={"facecolor": "white", "alpha": .8, "edgecolor": "none"})
    axis.set_ylabel("Normalized coupled residual")
    axis.grid(True, which="major", alpha=.22)
    axis.legend(fontsize=8, loc="best")
    title = "n = 3, J = 0: linear-solve tolerance control" if name == "n3" else "n = 2, J = 0: common-snapshot trial-path control"
    axis.set_title(title, pad=12)
    figure.supxlabel("Includes the initial norm at t = 0 and every completed Newton attempt; x marks a failed attempt.\n"
                      "Time excludes setup/observers. Residual convergence does not establish unbordered equilibrium or virial balance.",
                      fontsize=8)
    return figure, matplotlib.__version__


def sidecar(name, curves, sources):
    report = {
        "pair": name, "source_audits": sources,
        "plotter_sha256": digest_bytes(Path(__file__).read_bytes()),
        "curves": [{"id": c.identifier, "campaign": c.campaign, "audit_path": c.audit_path,
                    "producer_commit": c.producer_commit, "mesh_sha256": c.mesh_sha256,
                    "metadata": c.metadata, "input_state_hashes": c.input_state_hashes,
                    "cumulative_newton_seconds": c.time, "normalized_residual": c.residual,
                    "accepted_updates_after_t0": c.accepted, "actual_stopping_threshold": c.threshold,
                    "outcome": c.outcome, "failure": c.failure} for c in curves],
        "scope": ["Consumes the completed artifact-passed audit JSON; underlying file/FEM audits are not rerun.",
                  "Initial residual and stopping threshold are taken from the solver metric audited in that summary.",
                  "Cumulative time sums completed history rows; setup, observers and any pre-solver initialization are excluded.",
                  "Failed attempts retain their restored residual and elapsed time; no missing or incomplete case is drawn.",
                  "Zero residuals or thresholds are refused rather than silently floored on the logarithmic axis.",
                  "The n=3 norm may include the central-density border; neither residual decline nor its stopping threshold proves physical equilibrium or virial balance."],
    }
    if name == "n2":
        report["local_progress_panel"] = {
            "quantity": "100*(1-F/F0), in percent", "shared_initial_residual": curves[0].residual[0],
            "zero_included": True, "actual_stopping_threshold_display": "full-scale logarithmic panel above",
            "curves": [{"id": c.identifier, "residual_reduction_percent": residual_reduction(c)} for c in curves],
        }
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("audits", nargs="+", type=Path)
    parser.add_argument("--output", required=True, type=Path, help="New output directory outside the source checkout")
    parser.add_argument("--pair", choices=("both", "n3", "n2"), default="both")
    args = parser.parse_args()
    require(not args.output.exists(), "Refusing to overwrite an existing figure directory")
    require(not args.output.resolve().is_relative_to(Path(__file__).resolve().parents[1]),
            "Write plot artifacts outside the source checkout")
    pairs, sources = load_pairs(args.audits, args.pair)
    # Preflight every requested curve before importing Matplotlib or creating
    # output, so missing/incomplete controls cannot silently produce half a set.
    args.output.mkdir(parents=True)
    for name, curves in pairs.items():
        figure, version = render_pair(name, curves)
        report = sidecar(name, curves, sources)
        report["matplotlib_version"], report["python_version"] = version, sys.version
        report["figures"] = []
        for suffix in ("png", "pdf"):
            path = args.output / f"{name}_paired_convergence.{suffix}"
            figure.savefig(path, dpi=200)
            report["figures"].append({"path": str(path.resolve()), "sha256": digest_bytes(path.read_bytes())})
        from matplotlib import pyplot as plt
        plt.close(figure)
        path = args.output / f"{name}_paired_convergence.json"
        path.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n")
        print(path)


if __name__ == "__main__":
    main()