MeanField/experiments/summarize_density_order.py

#!/usr/bin/env python3
"""Summarize completed density-order audits without opening campaign artifacts.

The input is audit_density_order.py JSON with artifact_passed=true. Tables keep
zero, unavailable and uncertain observations distinct. Matplotlib is imported
only when a figure is requested; --self-test performs tiny arithmetic tests.
"""
from __future__ import annotations

import argparse
from collections import defaultdict
import csv
import hashlib
import json
import math
import os
from pathlib import Path
import unittest


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


def finite(value, label, nonnegative=False):
    require(isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value),
            "Missing/nonfinite " + label)
    require(not nonnegative or value >= 0, "Negative " + label)
    return value


def ratio(numerator, denominator):
    return numerator / denominator if numerator is not None and denominator not in (None, 0) else None


def density_metrics(case, key):
    measure = case["measures"][key]
    fine = measure["summary"]["30"]
    error = case["unsliced_error_quadrature"][key]
    q30, q26 = error["30"], error["26"]
    for field in ("negative_mass", "negative_volume", "negative_mass_fraction", "negative_volume_fraction",
                  "negative_mass_uncertainty_scale", "resolved_negative_intervals", "uncertain_negative_intervals"):
        finite(fine[field], field, True)
    squared = q30["eos_squared_integral"]
    if q30["eos_defined_everywhere_sampled"]:
        finite(squared, "q30 EOS squared integral", True)
        finite(q30["volume"], "q30 volume", True)
        require(q30["volume"] > 0, "Zero q30 volume")
        l2, rms = math.sqrt(squared), math.sqrt(squared / q30["volume"])
    else:
        require(squared is None, "Undefined q30 EOS integral has a fabricated value")
        l2 = rms = None
    if fine["negative_mass"] == 0:
        classification = "zero_observed"
    elif fine["resolved_negative_intervals"]:
        classification = "resolved_witnesses_present"
    else:
        classification = "uncertain_witnesses_only"
    surface_h = measure["surface"].get("enthalpy", {})
    witness_h = measure.get("negative_interval_witness_enthalpy", {})
    return {
        "negative_mass_over_abs_mass": fine["negative_mass_fraction"],
        "negative_volume_fraction": fine["negative_volume_fraction"],
        "negative_mass": fine["negative_mass"], "mass_denominator": abs(finite(fine["mass"], "mass")),
        "minimum_density_sliced": finite(fine["minimum_density"], "minimum density"),
        "negative_classification": classification,
        "resolved_negative_intervals": fine["resolved_negative_intervals"],
        "uncertain_negative_intervals": fine["uncertain_negative_intervals"],
        "uncertain_negative_mass": fine["uncertain_negative_mass"],
        "negative_mass_uncertainty_scale": fine["negative_mass_uncertainty_scale"],
        "negative_mass_over_uncertainty_scale": ratio(fine["negative_mass"], fine["negative_mass_uncertainty_scale"]),
        "negative_mass_q30_minus_q22": measure["negative_mass_q30_minus_q22"],
        "negative_volume_q30_minus_q22": measure["negative_volume_q30_minus_q22"],
        "eos_l2_q30": l2, "eos_rms_q30": rms,
        "eos_squared_q26": q26["eos_squared_integral"], "eos_squared_q30": squared,
        "eos_squared_q30_minus_q26": error["eos_squared_q30_minus_q26"],
        "eos_squared_relative_q30_minus_q26": ratio(error["eos_squared_q30_minus_q26"], squared),
        "q30_negative_enthalpy_integration_nodes": q30["negative_enthalpy_samples"],
        "surface_minimum_enthalpy": surface_h.get("minimum_enthalpy"),
        "surface_negative_enthalpy_samples": surface_h.get("negative_enthalpy_samples"),
        "interval_witness_minimum_enthalpy": witness_h.get("minimum_enthalpy"),
        "interval_negative_enthalpy_witnesses": witness_h.get("negative_enthalpy_samples")}


def build_rows(audits):
    """Deduplicate overlapping audited case records, never average repeat runs."""
    records = {}
    for audit_path, audit in audits:
        require(audit.get("artifact_passed") is True, "Refusing an incomplete/failed audit: " + audit_path)
        for campaign, run in audit["campaigns"].items():
            for name, case in run["cases"].items():
                require(case.get("artifact_passed") is True, "Case artifact audit failed: " + name)
                key = campaign, name
                payload = run["producer_commit"], run["mesh_sha256"], case
                if key in records:
                    require(records[key][1] == payload, "Overlapping audit records disagree: " + str(key))
                else:
                    records[key] = audit_path, payload
    frozen, coupled = [], []
    for (campaign, name), (audit_path, (commit, mesh, case)) in sorted(records.items()):
        common = {"audit_json": audit_path, "campaign": campaign, "case": name, "producer_commit": commit,
                  "mesh_sha256": mesh, "n": case["n"], "J": case["J"], "K": case["K"],
                  "source_density_degree": case["density_degree"], "source_state_dofs": case["state_size"]}
        if case["kind"] == "frozen":
            for degree in (2, 3, 4):
                key = f"Q{degree}/eos_l2_q26"
                row = {**common, "density_degree": degree, "control": "eos_l2_q26",
                       "density_dofs": case["projection"][f"Q{degree}"]["active_density_dofs"],
                       **density_metrics(case, key)}
                frozen.append(row)
        else:
            require(case["kind"] == "coupled", "Unknown case kind")
            newton = case["newton"]
            require(isinstance(newton["converged"], bool), "Invalid convergence observation")
            initial = finite(newton["initial_residual"], "initial residual", True)
            final = finite(newton["final_residual"], "final residual", True)
            reduction = ratio(final, initial)
            row = {**common, "density_degree": case["density_degree"], "control": "stored_density",
                   "density_dofs": case["block_sizes"]["value/density"], **density_metrics(case, f"Q{case['density_degree']}/stored_density"),
                   "converged": newton["converged"], "failure": newton["failure"],
                   "initial_residual": initial, "final_residual": final, "final_over_initial_residual": reduction,
                   "residual_reduction_percent": 100 * (1 - reduction) if reduction is not None else None,
                   "nonlinear_threshold": newton["threshold"], "accepted_steps": newton["accepted_steps"],
                   "iteration_records": len(newton["iterations"]), "krylov_iterations": newton["linear_iterations"],
                   "newton_iteration_seconds": finite(newton["iteration_seconds"], "Newton time", True),
                   "linear_solve_seconds": math.fsum(finite(r["linear_seconds"], "linear time", True) for r in newton["iterations"]),
                   "observer_seconds": case["observer_timings"]["observer_total"],
                   "elapsed_seconds_including_observers": case["run"]["seconds"],
                   "utc_minus_monotonic_seconds": case["run"]["utc_minus_monotonic_seconds"]}
            physical18 = case["physical"]["18"]["metrics"]
            # A bordered solve can meet its reported tolerance while the
            # unbordered equations still carry the auxiliary border action.
            for key in ("normalized_unbordered_residual", "central_border", "normalized_border_action"):
                row["final_" + key] = finite(physical18[key], "final " + key, nonnegative=key != "central_border")
            line_search = case["line_search"]
            for key in ("pressure_force_rejections", "residual_metric_rejections"):
                row[key] = finite(line_search[key], key, True)
            row["minimum_accepted_alpha"] = line_search["minimum_accepted_alpha"]
            if row["minimum_accepted_alpha"] is not None:
                require(0 < finite(row["minimum_accepted_alpha"], "minimum accepted alpha") <= 1, "Invalid accepted step")
            for q in (14, 18):
                physical = case["physical"][str(q)]["metrics"]
                for source, label in (("enthalpy_virial_error", "h_energy_virial"),
                                      ("enthalpy_force_virial_error", "h_force_virial"),
                                      ("virial_error", "density_energy_virial"),
                                      ("force_virial_error", "density_force_virial")):
                    value = physical[source]
                    if value is not None:
                        finite(value, source, True)
                    row[f"{label}_q{q}"] = value
                    row[f"{label}_q{q}_within_1e5"] = value <= 1e-5 if value is not None else None
            coupled.append(row)
    # Repeated same-model/same-degree runs require explicit selection upstream;
    # their averaging would conceal convergence outcomes and timing variability.
    seen = set()
    for row in coupled:
        key = row["case"], row["n"], row["J"], row["mesh_sha256"], row["density_degree"]
        require(key not in seen, "Duplicate coupled model/degree; select one audited run: " + str(key))
        seen.add(key)
    return frozen, coupled


def comparisons(frozen, coupled):
    result = []
    for kind, table in (("frozen", frozen), ("coupled", coupled)):
        groups = defaultdict(list)
        for row in table:
            key = (row["campaign"], row["case"]) if kind == "frozen" else (row["case"], row["n"], row["J"], row["mesh_sha256"])
            groups[key].append(row)
        for group in groups.values():
            baseline = next((r for r in group if r["density_degree"] == 2), None)
            if baseline is None:
                continue
            keys = ["negative_mass_over_abs_mass", "negative_volume_fraction", "eos_l2_q30"]
            if kind == "coupled":
                keys += ["newton_iteration_seconds", "krylov_iterations", "final_residual", "h_energy_virial_q18", "h_force_virial_q18"]
            for row in group:
                if row["density_degree"] != 2:
                    result.append({"kind": kind, "case": row["case"], "n": row["n"], "J": row["J"],
                                   "baseline_campaign": baseline["campaign"], "campaign": row["campaign"],
                                   "density_degree": row["density_degree"],
                                   **{key + "_over_Q2": ratio(row[key], baseline[key]) for key in keys}})
    return result


def write_csv(path, table, fields=None):
    fields = fields or list(dict.fromkeys(key for row in table for key in row))
    with path.open("w", newline="") as stream:
        if fields:
            writer = csv.DictWriter(stream, fieldnames=fields, extrasaction="ignore")
            writer.writeheader()
            writer.writerows(table)


def plot_figure(frozen, coupled, output):
    os.environ.setdefault("MPLCONFIGDIR", str(output / "matplotlib_cache"))
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False,
                         "savefig.dpi": 190, "pdf.fonttype": 42})
    row_count = bool(frozen) + 2 * bool(coupled)
    require(row_count > 0, "No audited rows to plot")
    fig, axes = plt.subplots(row_count, 3, figsize=(13.8, 3.9 * row_count + .7), squeeze=False)

    def groups(table, frozen_case):
        result = defaultdict(list)
        for row in table:
            key = (row["campaign"], row["case"]) if frozen_case else (row["case"], row["n"], row["J"], row["mesh_sha256"])
            result[key].append(row)
        return [sorted(group, key=lambda r: r["density_degree"]) for group in result.values()]

    def panel(ax, grouped, field, title, ylabel, negative=False, log=True, second=None, legend=False):
        annotations, positive = [], False
        for index, group in enumerate(grouped):
            color = plt.get_cmap("tab10")(index % 10)
            label = group[0]["case"]
            x = [r["density_degree"] for r in group]
            series = [(field, "-")] + ([(second, "--")] if second else [])
            for metric, style in series:
                values = []
                for row in group:
                    value = row[metric]
                    if value is None or (log and value == 0):
                        annotations.append(f"{label} Q{row['density_degree']}: {'undefined' if value is None else 'observed zero'}"
                                           + (" (force)" if metric == second else ""))
                        values.append(math.nan)
                        continue
                    finite(value, "plotted " + metric, log)
                    positive = True
                    values.append(value)
                    not_converged = row.get("converged") is False
                    uncertain = negative and row["negative_classification"] == "uncertain_witnesses_only"
                    ax.plot(row["density_degree"], value, marker="x" if not_converged else "o", markersize=5,
                            color=color, markerfacecolor="none" if uncertain else color, linestyle="none")
                ax.plot(x, values, style, color=color, linewidth=1.15, label=label if metric == field else None)
        if log and positive:
            ax.set_yscale("log")
        if not positive:
            ax.text(.5, .5, "No positive finite observations" if log else "No finite observations",
                    ha="center", va="center", transform=ax.transAxes)
        ax.set_xticks([2, 3, 4], ["Q2", "Q3", "Q4"])
        ax.set_xlim(1.85, 4.15)
        ax.set(title=title, ylabel=ylabel)
        ax.grid(alpha=.2, which="both")
        if annotations:
            ax.text(0, -.19, "\n".join(annotations), fontsize=6.5, transform=ax.transAxes, va="top")
        if legend:
            ax.legend(fontsize=7, loc="best")
        if second:
            handles = [Line2D([], [], color=".3", linestyle="-", label="Energy virial"),
                       Line2D([], [], color=".3", linestyle="--", label="Force virial")]
            ax.legend(handles=handles, fontsize=7)
            ax.axhline(1e-5, color=".4", linestyle=":", linewidth=.9)

    row = 0
    if frozen:
        grouped = groups(frozen, True)
        panel(axes[row, 0], grouped, "negative_mass_over_abs_mass", "Frozen L2 projection: negative mass", r"$M_- / |M|$", negative=True, legend=True)
        panel(axes[row, 1], grouped, "negative_volume_fraction", "Frozen L2 projection: negative volume", r"$V_- / V$", negative=True)
        panel(axes[row, 2], grouped, "eos_l2_q30", "Frozen L2 projection: EOS error", r"$\|\rho-\rho_{\rm EOS}(h)\|_{L^2}$ (q30)")
        row += 1
    if coupled:
        grouped = groups(coupled, False)
        panel(axes[row, 0], grouped, "negative_mass_over_abs_mass", "Coupled endpoint: negative mass", r"$M_- / |M|$", negative=True, legend=True)
        panel(axes[row, 1], grouped, "eos_l2_q30", "Coupled endpoint: EOS error", r"$\|\rho-\rho_{\rm EOS}(h)\|_{L^2}$ (q30)")
        panel(axes[row, 2], grouped, "h_energy_virial_q18", "Coupled endpoint: P(h) virials", "Absolute virial residual / |W| (q18)", second="h_force_virial_q18")
        row += 1
        panel(axes[row, 0], grouped, "final_over_initial_residual", "Reported nonlinear residual reduction", r"$\|F_{\rm final}\| / \|F_{\rm initial}\|$")
        panel(axes[row, 1], grouped, "newton_iteration_seconds", "Newton iteration cost", "Seconds (excludes observer and setup)", log=False)
        panel(axes[row, 2], grouped, "krylov_iterations", "Total reported Krylov work", "Iterations", log=False)
    fig.suptitle("Density approximation order — audited observations", fontsize=14, y=.985)
    fig.text(.025, .018, "Open markers: uncertain-only negative witnesses; ×: nonlinear solve did not converge. Zero/undefined values are listed below panels, not placed on a log floor.\n"
             "Negative regions use root-sliced quadrature; EOS norms use unsliced q30. P(h) virials do not certify density-EOS validity. See CSV/JSON for all outcomes and timing scope.", fontsize=8)
    fig.subplots_adjust(left=.075, right=.985, top=.92 if row_count == 1 else .95, bottom=.15 if row_count == 1 else .085,
                        wspace=.35, hspace=.72)
    fig.savefig(output / "density_order_comparison.png")
    fig.savefig(output / "density_order_comparison.pdf")
    plt.close(fig)


class Tests(unittest.TestCase):
    def test_zero_and_undefined_are_not_floored(self):
        fine = {k: 0. for k in ("negative_mass", "negative_volume", "negative_mass_fraction", "negative_volume_fraction",
                "negative_mass_uncertainty_scale", "resolved_negative_intervals", "uncertain_negative_intervals", "uncertain_negative_mass")}
        fine.update(mass=1., minimum_density=0.)
        measure = {"summary": {"30": fine}, "surface": {}, "negative_mass_q30_minus_q22": 0., "negative_volume_q30_minus_q22": 0.}
        q = {"eos_squared_integral": 0., "eos_defined_everywhere_sampled": True, "volume": 1., "negative_enthalpy_samples": 0}
        error = {"30": q, "26": q.copy(), "eos_squared_q30_minus_q26": 0.}
        case = {"measures": {"Q2/stored_density": measure}, "unsliced_error_quadrature": {"Q2/stored_density": error}}
        row = density_metrics(case, "Q2/stored_density")
        self.assertEqual(row["negative_classification"], "zero_observed")
        self.assertEqual(row["eos_l2_q30"], 0.)
        self.assertIsNone(row["negative_mass_over_uncertainty_scale"])
        q.update(eos_squared_integral=None, eos_defined_everywhere_sampled=False, negative_enthalpy_samples=1)
        self.assertIsNone(density_metrics(case, "Q2/stored_density")["eos_l2_q30"])
        fine.update(negative_mass=1e-20, negative_mass_fraction=1e-20, uncertain_negative_intervals=1)
        self.assertEqual(density_metrics(case, "Q2/stored_density")["negative_classification"], "uncertain_witnesses_only")

    def test_failed_audit_rejected(self):
        with self.assertRaises(ValueError):
            build_rows([("bad.json", {"artifact_passed": False})])

    def test_zero_baseline_ratio_undefined(self):
        self.assertIsNone(ratio(1., 0.))
        self.assertEqual(ratio(0., 1.), 0.)
        self.assertIsNone(ratio(None, 1.))

    def test_audited_coupled_outcome_keeps_iteration_cost(self):
        fine = {k: 0. for k in ("negative_mass", "negative_volume", "negative_mass_fraction", "negative_volume_fraction",
                "negative_mass_uncertainty_scale", "resolved_negative_intervals", "uncertain_negative_intervals", "uncertain_negative_mass")}
        fine.update(mass=1., minimum_density=0.)
        measure = {"summary": {"30": fine}, "surface": {}, "negative_mass_q30_minus_q22": 0., "negative_volume_q30_minus_q22": 0.}
        q = {"eos_squared_integral": 4e-8, "eos_defined_everywhere_sampled": True, "volume": 1., "negative_enthalpy_samples": 0}
        metrics = {"enthalpy_virial_error": 2e-6, "enthalpy_force_virial_error": 3e-5,
                   "normalized_unbordered_residual": .04, "central_border": -.125, "normalized_border_action": .03,
                   "virial_error": None, "force_virial_error": None}
        case = {"artifact_passed": True, "kind": "coupled", "n": 1.5, "J": 0., "K": 1., "density_degree": 2,
                "state_size": 50, "block_sizes": {"value/density": 27},
                "measures": {"Q2/stored_density": measure}, "unsliced_error_quadrature": {"Q2/stored_density": {
                    "30": q, "26": q, "eos_squared_q30_minus_q26": 0.}},
                "newton": {"converged": False, "failure": "iteration limit", "initial_residual": .1,
                    "final_residual": .01, "threshold": 1e-8, "accepted_steps": 1, "linear_iterations": 7,
                    "iteration_seconds": 10., "iterations": [{"linear_seconds": 6.}]},
                "observer_timings": {"observer_total": 200.}, "run": {"seconds": 215., "utc_minus_monotonic_seconds": .01},
                "line_search": {"pressure_force_rejections": 4, "residual_metric_rejections": 1, "minimum_accepted_alpha": 1/32},
                "physical": {"14": {"metrics": metrics}, "18": {"metrics": metrics}}}
        audit = {"artifact_passed": True, "campaigns": {"/recorded/campaign": {
            "producer_commit": "test", "mesh_sha256": "test", "cases": {"sample": case}}}}
        frozen, coupled = build_rows([("test.json", audit)])
        self.assertEqual(frozen, [])
        row = coupled[0]
        self.assertEqual(row["newton_iteration_seconds"], 10.)
        self.assertEqual(row["observer_seconds"], 200.)
        self.assertAlmostEqual(row["residual_reduction_percent"], 90.)
        self.assertFalse(row["converged"])
        self.assertTrue(row["h_energy_virial_q18_within_1e5"])
        self.assertFalse(row["h_force_virial_q18_within_1e5"])
        self.assertIsNone(row["density_energy_virial_q18_within_1e5"])
        self.assertEqual(row["final_normalized_unbordered_residual"], .04)
        self.assertEqual(row["final_central_border"], -.125)
        self.assertEqual(row["final_normalized_border_action"], .03)
        self.assertEqual(row["pressure_force_rejections"], 4)
        self.assertEqual(row["residual_metric_rejections"], 1)
        self.assertEqual(row["minimum_accepted_alpha"], 1/32)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("audits", nargs="*", type=Path)
    parser.add_argument("--output", type=Path, help="Fresh output directory")
    parser.add_argument("--tables-only", action="store_true")
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        test = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(Tests))
        if not test.wasSuccessful():
            raise SystemExit(1)
        if not args.audits:
            return
    if not args.audits or args.output is None:
        parser.error("Supply completed audit JSON and --output")
    inputs = [(str(path.resolve()), json.loads(path.read_text())) for path in args.audits]
    frozen, coupled = build_rows(inputs)
    require(frozen or coupled, "No cases in input audits")
    relative = comparisons(frozen, coupled)
    args.output.mkdir(parents=True, exist_ok=False)
    write_csv(args.output / "frozen_projection.csv", frozen, ["case", "n", "J", "density_degree", "density_dofs",
        "negative_mass_over_abs_mass", "negative_volume_fraction", "negative_classification", "negative_mass_over_uncertainty_scale",
        "eos_l2_q30", "eos_rms_q30", "eos_squared_relative_q30_minus_q26", "negative_mass_q30_minus_q22",
        "surface_minimum_enthalpy", "surface_negative_enthalpy_samples"])
    write_csv(args.output / "coupled_outcomes.csv", coupled, ["case", "n", "J", "density_degree", "source_state_dofs", "density_dofs",
        "converged", "failure", "accepted_steps", "krylov_iterations", "newton_iteration_seconds", "linear_solve_seconds",
        "pressure_force_rejections", "residual_metric_rejections", "minimum_accepted_alpha",
        "initial_residual", "final_residual", "final_over_initial_residual", "h_energy_virial_q18", "h_force_virial_q18",
        "final_normalized_unbordered_residual", "final_central_border", "final_normalized_border_action",
        "h_energy_virial_q14", "h_force_virial_q14", "negative_mass_over_abs_mass", "negative_volume_fraction",
        "negative_classification", "eos_l2_q30", "eos_rms_q30"])
    write_csv(args.output / "relative_to_Q2.csv", relative)
    if not args.tables_only:
        plot_figure(frozen, coupled, args.output)
    digest = lambda path: hashlib.sha256(path.read_bytes()).hexdigest()
    summary = {"schema": "density_order_summary_v1", "generator_sha256": digest(Path(__file__)),
               "input_audit_sha256": {str(path.resolve()): digest(path) for path in args.audits},
               "scope": "Presentation of audited observations only; no campaign files reopened and no independent FEM work.",
               "notes": ["Frozen controls are fixed-q26 L2 projections; EOS error norms use separate unsliced q30.",
                         "Negative-mass denominator is the absolute measured density mass, not a substituted exact target mass.",
                         "Uncertainty scales are reconstruction indicators, not rigorous bounds; zero observed negativity is not a global positivity certificate.",
                         "Newton seconds sum reported nonlinear-iteration times, excluding observer and setup costs; total elapsed time is kept separately.",
                         "Final unbordered residual, signed central border and normalized border action come from q18 physical diagnostics; bordered convergence does not imply that these vanish.",
                         "Pressure-force and residual-metric rejection counts retain the recorded stage distinction; residual-metric counts include any nonfinite metric failures as well as insufficient decrease.",
                         "Ratios to Q2 compare recorded endpoints and do not imply equal convergence or error tolerances were attained.",
                         "Undefined values are null/blank; zero is kept exactly. No logarithmic plot floor is inserted.",
                         "P(h) virial milestones do not establish pointwise density validity; density-based virials are retained separately."],
               "frozen_projection": frozen, "coupled_outcomes": coupled, "relative_to_Q2": relative,
               "artifacts_sha256": {path.name: digest(path) for path in args.output.iterdir() if path.is_file()}}
    (args.output / "summary.json").write_text(json.dumps(summary, indent=2, allow_nan=False) + "\n")
    print("Wrote", args.output)


if __name__ == "__main__":
    main()