MeanField/experiments/compare_model_slices.py

#!/usr/bin/env python3
"""Compare two visualization resolutions of the same physical FEM state."""
from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np

from plot_model_slices import Slice, color_scale, load_slice, plt, render


def slice_metrics(data: Slice) -> dict:
    """Integrate the displayed triangles, not the underlying curved FE slice."""
    u, v = int(data.metadata["horizontal_axis"]), int(data.metadata["vertical_axis"])
    points = data.vertices[data.triangles]
    x, y = points[:, :, 2 + u], points[:, :, 2 + v]
    area = .5 * np.abs((x[:, 1] - x[:, 0]) * (y[:, 2] - y[:, 0]) -
                       (x[:, 2] - x[:, 0]) * (y[:, 1] - y[:, 0]))
    metrics = {
        "triangle_count": len(data.triangles),
        "triangle_area_sum": float(np.sum(area)),
        "piecewise_linear_density_integral": float(np.sum(area * points[:, :, 5].mean(axis=1))),
        "piecewise_linear_enthalpy_integral": float(np.sum(area * points[:, :, 6].mean(axis=1))),
    }
    metrics["exported_geometric_diagnostics"] = {
        key: float(data.metadata[key]) if key in data.metadata else None
        for key in ("maximum_plane_error", "maximum_sampled_chord_error", "surface_maximum_plane_error",
                    "reference_surface_maximum_plane_error", "surface_maximum_sampled_chord_error",
                    "reference_surface_maximum_sampled_chord_error")
    }
    for name, segments in (("surface", data.surface), ("reference_surface", data.reference_surface)):
        if len(segments):
            a, b = segments[:, [u, v]], segments[:, [3 + u, 3 + v]]
            endpoints = np.concatenate((a, b))
            values = {"segment_count": len(segments),
                      "segment_length_sum": float(np.linalg.norm(b - a, axis=1).sum()),
                      "u_min": float(endpoints[:, 0].min()), "u_max": float(endpoints[:, 0].max()),
                      "v_min": float(endpoints[:, 1].min()), "v_max": float(endpoints[:, 1].max())}
        else:
            values = {"segment_count": 0, "segment_length_sum": 0., "u_min": None, "u_max": None, "v_min": None, "v_max": None}
        metrics[name] = values
    return metrics


def differences(coarse: dict, fine: dict) -> dict:
    result = {}
    for key in coarse:
        a, b = coarse[key], fine[key]
        if isinstance(a, dict):
            result[key] = differences(a, b)
        elif a is None or b is None:
            result[key] = {"coarse": a, "fine": b, "signed_change": None, "absolute_change": None,
                           "relative_change": None, "relative_absolute_change": None}
        else:
            delta = b - a
            relative = delta / abs(a) if a != 0 else (0. if delta == 0 else None)
            result[key] = {"coarse": a, "fine": b, "signed_change": delta, "absolute_change": abs(delta),
                           "relative_change": relative,
                           "relative_absolute_change": abs(relative) if relative is not None else None}
    return result


def validate_pair(coarse: Slice, fine: Slice) -> None:
    if coarse.directory.resolve() == fine.directory.resolve():
        raise ValueError("Coarse and fine exports must be distinct directories")
    for key in ("plane", "horizontal_axis", "vertical_axis", "normal_axis"):
        if coarse.metadata[key] != fine.metadata[key]:
            raise ValueError("The paired exports have different physical planes")
    if float(coarse.metadata["offset"]) != float(fine.metadata["offset"]):
        raise ValueError("The paired exports have different physical plane offsets")
    if int(coarse.metadata["subdivisions"]) >= int(fine.metadata["subdivisions"]):
        raise ValueError("Fine subdivisions must exceed coarse subdivisions")


def comparison(coarse: Slice, fine: Slice, source_state: Path) -> dict:
    validate_pair(coarse, fine)
    if not source_state.is_file():
        raise ValueError("--source-state must identify the caller-verified common saved state file")
    digest = hashlib.sha256()
    with source_state.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return {
        "coarse_directory": str(coarse.directory.resolve()), "fine_directory": str(fine.directory.resolve()),
        "plane": coarse.metadata["plane"], "offset": float(coarse.metadata["offset"]),
        "coarse_subdivisions": int(coarse.metadata["subdivisions"]), "fine_subdivisions": int(fine.metadata["subdivisions"]),
        "state_provenance": {
            "caller_asserted_common_source": str(source_state.resolve()), "source_sha256": digest.hexdigest(),
            "export_binding_verified": False,
            "limitation": "Caller must verify that both exports used this same state, mesh, field orders and physical specifications. The CSVs do not bind themselves cryptographically to the source state.",
        },
        "changes": differences(slice_metrics(coarse), slice_metrics(fine)),
        "scope": [
            "Triangle area and field integrals sum supplied triangles, including any overlaps; they are not union-area measures.",
            "Field integrals are signed integrals of piecewise-linear vertex data in 2D; they are not 3D mass, pressure, virial or exact curved-FEM integrals.",
            "Surface segment-length sums approximate perimeter and include any duplicated/overlapping segments.",
            "Relative change is (fine-coarse)/abs(coarse); it is null for a zero coarse value and a nonzero fine value.",
            "Near-zero coarse signed integrals can have large relative changes; inspect absolute changes as well.",
            "Shared physical axes and separate shared density/enthalpy color normalizations preserve actual deformation and signed field values.",
            "Agreement under one refinement does not certify global geometric conformity, positivity or equilibrium.",
        ],
    }


def render_comparison(coarse: Slice, fine: Slice, *, reference_outline=False, title=None, pixels=512):
    validate_pair(coarse, fine)
    bounds = (min(float(d.metadata["u_min"]) for d in (coarse, fine)),
              max(float(d.metadata["u_max"]) for d in (coarse, fine)),
              min(float(d.metadata["v_min"]) for d in (coarse, fine)),
              max(float(d.metadata["v_max"]) for d in (coarse, fine)))
    normalizations = [color_scale(np.concatenate((coarse.vertices[:, column], fine.vertices[:, column])))
                      for column in (5, 6)] if len(coarse.vertices) + len(fine.vertices) else None
    figure, axes = plt.subplots(2, 2, figsize=(13.4, 10.6), layout="constrained")
    for column, (label, data) in enumerate((("Coarse", coarse), ("Fine", fine))):
        render(data, reference_outline=reference_outline,
               title=f"{label}: {data.metadata['subdivisions']} subdivisions",
               normalizations=normalizations, physical_bounds=bounds,
               figure=figure, axes=axes[:, column], heading=False, pixels=pixels)
    figure.suptitle(f"{title or 'Same-state slice refinement'} — physical {coarse.metadata['plane']} plane, offset {float(coarse.metadata['offset']):.6g}\n"
                   f"Common axes/color scales; {pixels}×{pixels} scalar pixels; actual deformation; equilibrium not assessed", fontsize=11)
    return figure


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("coarse", type=Path, help="Coarse plane directory, e.g. first/final_slices/xz")
    parser.add_argument("fine", type=Path, help="Fine plane directory for the same saved state")
    parser.add_argument("--source-state", required=True, type=Path,
                        help="Common saved state file; caller asserts both exports used this state and matching discretization")
    parser.add_argument("--output", required=True, type=Path, help="New JSON/PNG/PDF prefix outside the source checkout")
    parser.add_argument("--reference-outline", action="store_true")
    parser.add_argument("--pixels", type=int, default=512, help="Common scalar sample grid per axis, 32–2048")
    parser.add_argument("--title")
    args = parser.parse_args()
    source_checkout = Path(__file__).resolve().parents[1]
    if args.output.resolve().is_relative_to(source_checkout):
        raise ValueError("Comparison artifacts must be written outside the source checkout")
    paths = [args.output.with_suffix(suffix) for suffix in (".json", ".png", ".pdf")]
    if any(path.exists() for path in paths):
        raise FileExistsError("Refusing to overwrite an existing comparison artifact")
    coarse, fine = load_slice(args.coarse), load_slice(args.fine)
    report = comparison(coarse, fine, args.source_state)
    figure = render_comparison(coarse, fine, reference_outline=args.reference_outline, title=args.title, pixels=args.pixels)
    report["rasterization"] = {"pixels_per_axis": args.pixels, "sample_location": "pixel centers",
                              "scalar_interpolation": "barycentric within each supplied triangle before color mapping",
                              "overlap_owner": "first supplied triangle; no trace averaging", "vacuum": "masked",
                              "limitation": "Subpixel features may be missed; this raster is not an exact curved-FEM evaluation."}
    report["figures"] = [str(path.resolve()) for path in paths[1:]]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    paths[0].write_text(json.dumps(report, indent=2, allow_nan=False) + "\n")
    for path in paths[1:]:
        figure.savefig(path, dpi=200, bbox_inches="tight", pad_inches=.15)
    plt.close(figure)
    for path in paths:
        print(path)


if __name__ == "__main__":
    main()