drbx 2.0.0.dev0__py3-none-any.whl

drbx/validation/fluid_1d_mms_convergence.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+from matplotlib import pyplot as plt
+import numpy as np
+from .publication_plotting import annotate_bars, save_publication_figure, style_axis
+
+from ..config.boutinp import parse_bout_input
+from ..native.fluid_1d import advance_mms_history, evaluate_field_option
+from ..native.mesh import build_structured_mesh
+from ..native.metrics import build_structured_metrics
+from ..runtime.run_config import RunConfiguration
+
+
+_MMS_TEMPLATE = """
+nout = 50
+timestep = {timestep}
+
+MXG = 0
+
+[mesh]
+nx = 1
+ny = {ny}
+nz = 1
+Ly = 10
+dy = Ly / ny
+J = 1
+
+[solver]
+mxstep = 10000
+rtol = 1e-7
+mms = true
+
+[model]
+components = i
+normalise_metric = false
+Nnorm = 1e18
+Bnorm = 1
+Tnorm = 5
+
+[i]
+type = evolve_density, evolve_pressure, evolve_momentum
+charge = 1.0
+AA = 2.0
+thermal_conduction = false
+
+[Ni]
+solution = 1 - 0.1*sin(t - 2.0*y)
+source = -0.1*cos(t - 2.0*y) + 0.0628318530717959*cos(2*t + y)
+
+[Pi]
+solution = 0.1*cos(t + 3.0*y) + 1
+source = (0.0628318530717959*cos(2*t + y)/(1 - 0.1*sin(t - 2.0*y)) - 0.0125663706143592*sin(2*t + y)*cos(t - 2.0*y)/(1 - 0.1*sin(t - 2.0*y))^2)*(0.0666666666666667*cos(t + 3.0*y) + 0.666666666666667) - 0.1*sin(t + 3.0*y) + 0.0628318530717959*(0.1*cos(t + 3.0*y) + 1)*cos(2*t + y)/(1 - 0.1*sin(t - 2.0*y)) - 0.0188495559215388*sin(t + 3.0*y)*sin(2*t + y)/(1 - 0.1*sin(t - 2.0*y)) - 0.0125663706143592*(0.1*cos(t + 3.0*y) + 1)*sin(2*t + y)*cos(t - 2.0*y)/(1 - 0.1*sin(t - 2.0*y))^2
+
+[NVi]
+solution = 0.2*sin(2*t + y)
+source = -0.188495559215388*sin(t + 3.0*y) + 0.4*cos(2*t + y) + 0.0251327412287183*sin(2*t + y)*cos(2*t + y)/(1 - 0.1*sin(t - 2.0*y)) - 0.00251327412287184*sin(2*t + y)^2*cos(t - 2.0*y)/(1 - 0.1*sin(t - 2.0*y))^2
+"""
+
+
+@dataclass(frozen=True)
+class Fluid1DMmsConvergenceArtifacts:
+    summary_json_path: Path
+    arrays_npz_path: Path
+    summary_plot_png_path: Path
+
+
+def create_fluid_1d_mms_convergence_package(
+    *,
+    output_root: str | Path,
+    case_label: str = "fluid_1d_mms_convergence",
+    resolutions: tuple[int, ...] = (32, 64, 128),
+    timestep: float = 0.05,
+    steps: int = 2,
+    substeps: int = 20,
+) -> Fluid1DMmsConvergenceArtifacts:
+    root = Path(output_root)
+    data_dir = root / "data"
+    images_dir = root / "images"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    report = build_fluid_1d_mms_convergence_report(
+        resolutions=resolutions,
+        timestep=timestep,
+        steps=steps,
+        substeps=substeps,
+    )
+    summary_json_path = data_dir / f"{case_label}.json"
+    summary_json_path.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
+
+    arrays_npz_path = data_dir / f"{case_label}.npz"
+    np.savez_compressed(
+        arrays_npz_path,
+        resolutions=np.asarray(report["resolutions"], dtype=np.int64),
+        density_l2=np.asarray([entry["errors"]["density_l2"] for entry in report["runs"]], dtype=np.float64),
+        pressure_l2=np.asarray([entry["errors"]["pressure_l2"] for entry in report["runs"]], dtype=np.float64),
+        momentum_l2=np.asarray([entry["errors"]["momentum_l2"] for entry in report["runs"]], dtype=np.float64),
+        density_order=np.asarray([entry["density_order"] for entry in report["observed_orders"]], dtype=np.float64),
+        pressure_order=np.asarray([entry["pressure_order"] for entry in report["observed_orders"]], dtype=np.float64),
+        momentum_order=np.asarray([entry["momentum_order"] for entry in report["observed_orders"]], dtype=np.float64),
+    )
+
+    summary_plot_png_path = save_fluid_1d_mms_convergence_plot(report, images_dir / f"{case_label}.png")
+    return Fluid1DMmsConvergenceArtifacts(
+        summary_json_path=summary_json_path,
+        arrays_npz_path=arrays_npz_path,
+        summary_plot_png_path=summary_plot_png_path,
+    )
+
+
+def build_mms_config(*, ny: int, timestep: float) -> str:
+    return _MMS_TEMPLATE.format(ny=int(ny), timestep=float(timestep))
+
+
+def l2_error(numerical: np.ndarray, exact: np.ndarray, *, start: int, end: int) -> float:
+    interior_numerical = np.asarray(numerical[:, start : end + 1, :], dtype=np.float64)
+    interior_exact = np.asarray(exact[:, start : end + 1, :], dtype=np.float64)
+    return float(np.sqrt(np.mean(np.square(interior_numerical - interior_exact))))
+
+
+def observed_order(coarse_error: float, fine_error: float) -> float:
+    if coarse_error <= 0.0 or fine_error <= 0.0:
+        return float("inf")
+    return float(np.log(coarse_error / fine_error) / np.log(2.0))
+
+
+def run_mms_resolution(*, ny: int, timestep: float, steps: int, substeps: int) -> dict[str, object]:
+    config = parse_bout_input(build_mms_config(ny=ny, timestep=timestep))
+    run_config = RunConfiguration.from_config(config)
+    mesh = build_structured_mesh(config, run_config)
+    metrics = build_structured_metrics(config, run_config, mesh)
+    history = advance_mms_history(
+        config,
+        section="i",
+        mesh=mesh,
+        metrics=metrics,
+        atomic_mass=2.0,
+        timestep=timestep,
+        steps=steps,
+        substeps=substeps,
+    )
+    final_time = float(steps * timestep)
+    exact_density = evaluate_field_option(config, "Ni", "solution", mesh=mesh, time=final_time)
+    exact_pressure = evaluate_field_option(config, "Pi", "solution", mesh=mesh, time=final_time)
+    exact_momentum = evaluate_field_option(config, "NVi", "solution", mesh=mesh, time=final_time)
+
+    density_error = l2_error(np.asarray(history.density_history[-1]), np.asarray(exact_density), start=mesh.ystart, end=mesh.yend)
+    pressure_error = l2_error(np.asarray(history.pressure_history[-1]), np.asarray(exact_pressure), start=mesh.ystart, end=mesh.yend)
+    momentum_error = l2_error(np.asarray(history.momentum_history[-1]), np.asarray(exact_momentum), start=mesh.ystart, end=mesh.yend)
+
+    return {
+        "ny": int(ny),
+        "timestep": float(timestep),
+        "steps": int(steps),
+        "substeps": int(substeps),
+        "final_time": final_time,
+        "errors": {
+            "density_l2": density_error,
+            "pressure_l2": pressure_error,
+            "momentum_l2": momentum_error,
+        },
+    }
+
+
+def build_fluid_1d_mms_convergence_report(
+    *,
+    resolutions: tuple[int, ...] = (32, 64, 128),
+    timestep: float = 0.05,
+    steps: int = 2,
+    substeps: int = 20,
+) -> dict[str, object]:
+    runs = [run_mms_resolution(ny=ny, timestep=timestep, steps=steps, substeps=substeps) for ny in resolutions]
+    orders: list[dict[str, float | int]] = []
+    for coarse, fine in zip(runs[:-1], runs[1:], strict=False):
+        coarse_errors = coarse["errors"]
+        fine_errors = fine["errors"]
+        orders.append(
+            {
+                "from_ny": int(coarse["ny"]),
+                "to_ny": int(fine["ny"]),
+                "density_order": observed_order(float(coarse_errors["density_l2"]), float(fine_errors["density_l2"])),
+                "pressure_order": observed_order(float(coarse_errors["pressure_l2"]), float(fine_errors["pressure_l2"])),
+                "momentum_order": observed_order(float(coarse_errors["momentum_l2"]), float(fine_errors["momentum_l2"])),
+            }
+        )
+    return {
+        "family": "manufactured_solution_convergence",
+        "case": "fluid_1d_mms_convergence",
+        "operator_family": "fluid_1d_density_pressure_momentum",
+        "literature_anchor": "Standard code-verification literature uses refinement studies and observed-order plots to separate operator correctness from benchmark validation.",
+        "resolutions": [int(ny) for ny in resolutions],
+        "timestep": float(timestep),
+        "steps": int(steps),
+        "substeps": int(substeps),
+        "runs": runs,
+        "observed_orders": orders,
+        "min_observed_order": {
+            "density": float(min(entry["density_order"] for entry in orders)),
+            "pressure": float(min(entry["pressure_order"] for entry in orders)),
+            "momentum": float(min(entry["momentum_order"] for entry in orders)),
+        },
+    }
+
+
+def save_fluid_1d_mms_convergence_plot(report: dict[str, object], path: str | Path) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    runs = list(report["runs"])
+    resolutions = np.asarray([entry["ny"] for entry in runs], dtype=np.float64)
+    density_error = np.asarray([entry["errors"]["density_l2"] for entry in runs], dtype=np.float64)
+    pressure_error = np.asarray([entry["errors"]["pressure_l2"] for entry in runs], dtype=np.float64)
+    momentum_error = np.asarray([entry["errors"]["momentum_l2"] for entry in runs], dtype=np.float64)
+    orders = list(report["observed_orders"])
+
+    figure, axes = plt.subplots(1, 2, figsize=(13.5, 5.2), constrained_layout=True)
+
+    axes[0].loglog(resolutions, density_error, marker="o", linewidth=2.0, color="#005f73", label="density")
+    axes[0].loglog(resolutions, pressure_error, marker="o", linewidth=2.0, color="#ca6702", label="pressure")
+    axes[0].loglog(resolutions, momentum_error, marker="o", linewidth=2.0, color="#3a86ff", label="momentum")
+    style_axis(
+        axes[0],
+        title="Fluid 1D MMS refinement errors",
+        xlabel="interior Ny resolution",
+        ylabel="L2 error",
+        xscale="log",
+        yscale="log",
+        grid="both",
+    )
+    axes[0].legend(frameon=False)
+
+    order_labels = [f"{entry['from_ny']}→{entry['to_ny']}" for entry in orders]
+    x = np.arange(len(order_labels))
+    width = 0.25
+    axes[1].bar(x - width, [entry["density_order"] for entry in orders], width=width, color="#005f73", label="density")
+    axes[1].bar(x, [entry["pressure_order"] for entry in orders], width=width, color="#ca6702", label="pressure")
+    axes[1].bar(x + width, [entry["momentum_order"] for entry in orders], width=width, color="#3a86ff", label="momentum")
+    axes[1].axhline(2.0, color="#bb3e03", linestyle="--", linewidth=1.5, label="second order")
+    axes[1].set_xticks(x, order_labels)
+    style_axis(axes[1], title="Observed MMS refinement order", ylabel="observed order")
+    axes[1].legend(frameon=False)
+    annotate_bars(axes[1], x - width, np.asarray([entry["density_order"] for entry in orders], dtype=np.float64), fmt="{:.2f}", fontsize=8.5)
+    annotate_bars(axes[1], x, np.asarray([entry["pressure_order"] for entry in orders], dtype=np.float64), fmt="{:.2f}", fontsize=8.5)
+    annotate_bars(axes[1], x + width, np.asarray([entry["momentum_order"] for entry in orders], dtype=np.float64), fmt="{:.2f}", fontsize=8.5)
+    figure.suptitle("Fluid 1D manufactured-solution convergence audit", fontsize=14.0, fontweight="semibold")
+    save_publication_figure(figure, target)
+    return target

drbx/validation/geometry_lineouts.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+import json
+
+from matplotlib import pyplot as plt
+from matplotlib.ticker import ScalarFormatter
+import numpy as np
+
+
+@dataclass(frozen=True)
+class LineoutSpec:
+    name: str
+    axis: int
+    coordinate_name: str
+    coordinate_values: np.ndarray
+    fixed_indices: tuple[int, int]
+
+
+def build_lineout_report(
+    *,
+    fields: dict[str, np.ndarray],
+    specs: tuple[LineoutSpec, ...],
+) -> dict[str, object]:
+    diagnostics: dict[str, dict[str, object]] = {}
+    for spec in specs:
+        diagnostics[spec.name] = {}
+        for field_name, values in fields.items():
+            line = _extract_line(values, spec)
+            diagnostics[spec.name][field_name] = {
+                "coordinate_name": spec.coordinate_name,
+                "coordinate_values": np.asarray(spec.coordinate_values, dtype=np.float64).tolist(),
+                "mean": line.tolist(),
+                "minimum": float(np.min(line)),
+                "maximum": float(np.max(line)),
+            }
+    return {"available": True, "parse_status": "ok", "diagnostics": diagnostics}
+
+
+def write_lineout_arrays_npz(report: dict[str, object], path: str | Path) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    payload: dict[str, np.ndarray] = {}
+    diagnostics = report.get("diagnostics", {})
+    if isinstance(diagnostics, dict):
+        for diagnostic_name, fields in diagnostics.items():
+            if not isinstance(fields, dict):
+                continue
+            for field_name, field_report in fields.items():
+                if not isinstance(field_report, dict):
+                    continue
+                key_prefix = f"{diagnostic_name}:{field_name}"
+                payload[f"{key_prefix}:coords"] = np.asarray(field_report.get("coordinate_values", []), dtype=np.float64)
+                payload[f"{key_prefix}:mean"] = np.asarray(field_report.get("mean", []), dtype=np.float64)
+    payload["__metadata__"] = np.asarray(json.dumps(report, sort_keys=True), dtype=np.str_)
+    np.savez_compressed(target, **payload)
+    return target
+
+
+def save_lineout_summary_plot(
+    report: dict[str, object],
+    path: str | Path,
+    *,
+    field_names: tuple[str, ...],
+    title: str,
+) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    diagnostics = report.get("diagnostics", {})
+    diagnostic_names = tuple(diagnostics.keys()) if isinstance(diagnostics, dict) else ()
+    figure, axes = plt.subplots(
+        len(field_names),
+        max(1, len(diagnostic_names)),
+        figsize=(13.5, max(6.0, 2.35 * len(field_names))),
+        constrained_layout=True,
+        squeeze=False,
+    )
+    for col, diagnostic_name in enumerate(diagnostic_names):
+        diagnostic = diagnostics.get(diagnostic_name, {})
+        for row, field_name in enumerate(field_names):
+            axis = axes[row, col]
+            field_report = diagnostic.get(field_name) if isinstance(diagnostic, dict) else None
+            if not isinstance(field_report, dict):
+                axis.set_visible(False)
+                continue
+            coords = np.asarray(field_report.get("coordinate_values", []), dtype=np.float64)
+            line = np.asarray(field_report.get("mean", []), dtype=np.float64)
+            axis.plot(coords, line, color="#005f73", linewidth=2.2)
+            axis.grid(alpha=0.25)
+            axis.set_title(f"{_display_label(diagnostic_name)} · {_display_label(field_name)}", fontsize=10)
+            axis.set_xlabel(str(field_report.get("coordinate_name", "coord")))
+            axis.set_ylabel(_display_label(field_name))
+            span = float(np.ptp(line))
+            scale = max(float(np.max(np.abs(line))), 1.0e-12)
+            formatter = ScalarFormatter(useOffset=False)
+            formatter.set_powerlimits((-2, 2))
+            if span / scale < 1.0e-3:
+                formatter.set_scientific(False)
+            axis.yaxis.set_major_formatter(formatter)
+    figure.suptitle(title, fontsize=16, fontweight="bold")
+    figure.savefig(target, dpi=180)
+    plt.close(figure)
+    return target
+
+
+def _extract_line(values: np.ndarray, spec: LineoutSpec) -> np.ndarray:
+    array = np.asarray(values, dtype=np.float64)
+    if array.ndim != 3:
+        raise ValueError(f"Expected 3D array for lineout extraction, got shape {array.shape}")
+    if spec.axis == 0:
+        return array[:, spec.fixed_indices[0], spec.fixed_indices[1]]
+    if spec.axis == 1:
+        return array[spec.fixed_indices[0], :, spec.fixed_indices[1]]
+    if spec.axis == 2:
+        return array[spec.fixed_indices[0], spec.fixed_indices[1], :]
+    raise ValueError(f"Unsupported axis {spec.axis}")
+
+
+def _display_label(name: str) -> str:
+    labels = {
+        "radial_midplane": "Radial midplane",
+        "toroidal_cut": "Toroidal cut",
+        "poloidal_cut": "Poloidal cut",
+        "Bxy": "B",
+        "Bmag": "|B|",
+        "J": "Jacobian",
+        "jacobian": "Jacobian",
+        "g11": "g11",
+        "g22": "g22",
+        "g33": "g33",
+        "g_11": "g^11",
+        "g_22": "g^22",
+        "g_33": "g^33",
+    }
+    return labels.get(name, name.replace("_", " "))

drbx/validation/geometry_slices.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+import json
+
+from matplotlib import pyplot as plt
+from matplotlib import animation
+import numpy as np
+
+
+@dataclass(frozen=True)
+class SliceSpec:
+    name: str
+    axis: int
+    coordinate_name: str
+    coordinate_values: np.ndarray
+
+
+def build_slice_report(
+    *,
+    field_name: str,
+    values: np.ndarray,
+    spec: SliceSpec,
+) -> dict[str, object]:
+    array = np.asarray(values, dtype=np.float64)
+    if array.ndim != 3:
+        raise ValueError(f"Expected 3D array for slice diagnostics, got shape {array.shape}")
+    frames = []
+    reference_plane = _extract_plane(array, spec.axis, 0)
+    for index, coord in enumerate(np.asarray(spec.coordinate_values, dtype=np.float64)):
+        plane = _extract_plane(array, spec.axis, index)
+        delta = plane - reference_plane
+        frames.append(
+            {
+                "index": int(index),
+                "coordinate_value": float(coord),
+                "minimum": float(np.min(plane)),
+                "maximum": float(np.max(plane)),
+                "mean": float(np.mean(plane)),
+                "standard_deviation": float(np.std(plane)),
+                "rms_delta_from_first": float(np.sqrt(np.mean(delta**2))),
+            }
+        )
+    return {
+        "available": True,
+        "parse_status": "ok",
+        "field_name": field_name,
+        "slice_name": spec.name,
+        "coordinate_name": spec.coordinate_name,
+        "coordinate_values": np.asarray(spec.coordinate_values, dtype=np.float64).tolist(),
+        "frames": frames,
+    }
+
+
+def write_slice_arrays_npz(
+    *,
+    field_name: str,
+    values: np.ndarray,
+    spec: SliceSpec,
+    path: str | Path,
+) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    array = np.asarray(values, dtype=np.float64)
+    payload = {
+        "field_name": np.asarray(field_name),
+        "coordinate_name": np.asarray(spec.coordinate_name),
+        "coordinate_values": np.asarray(spec.coordinate_values, dtype=np.float64),
+        "values": array,
+    }
+    np.savez_compressed(target, **payload)
+    return target
+
+
+def save_slice_summary_plot(
+    report: dict[str, object],
+    path: str | Path,
+    *,
+    title: str,
+) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    frames = report.get("frames", [])
+    coords = np.asarray([frame["coordinate_value"] for frame in frames], dtype=np.float64)
+    mins = np.asarray([frame["minimum"] for frame in frames], dtype=np.float64)
+    maxs = np.asarray([frame["maximum"] for frame in frames], dtype=np.float64)
+    means = np.asarray([frame["mean"] for frame in frames], dtype=np.float64)
+    stds = np.asarray([frame.get("standard_deviation", 0.0) for frame in frames], dtype=np.float64)
+    rms_deltas = np.asarray([frame.get("rms_delta_from_first", 0.0) for frame in frames], dtype=np.float64)
+    figure, axes = plt.subplots(2, 1, figsize=(10.5, 7.2), constrained_layout=True, sharex=True)
+    axes[0].plot(coords, means, color="#005f73", linewidth=2.0, label="mean")
+    axes[0].fill_between(coords, mins, maxs, color="#94d2bd", alpha=0.35, label="min/max band")
+    axes[0].set_ylabel(_display_label(str(report.get("field_name", "field"))))
+    axes[0].set_title(title)
+    axes[0].grid(alpha=0.25)
+    axes[0].legend(frameon=False)
+    axes[0].ticklabel_format(axis="y", style="sci", scilimits=(-2, 2), useOffset=False)
+    axes[1].plot(coords, stds, color="#0a9396", linewidth=1.9, label="plane std")
+    axes[1].plot(coords, rms_deltas, color="#ae2012", linewidth=1.9, label="RMS delta from first")
+    axes[1].set_xlabel(str(report.get("coordinate_name", "coord")))
+    axes[1].set_ylabel("diagnostic amplitude")
+    axes[1].grid(alpha=0.25)
+    axes[1].legend(frameon=False)
+    axes[1].ticklabel_format(axis="y", style="sci", scilimits=(-2, 2), useOffset=False)
+    figure.savefig(target, dpi=180)
+    plt.close(figure)
+    return target
+
+
+def save_slice_gif(
+    *,
+    field_name: str,
+    values: np.ndarray,
+    spec: SliceSpec,
+    path: str | Path,
+    fps: int = 8,
+) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    array = np.asarray(values, dtype=np.float64)
+    frames = [_extract_plane(array, spec.axis, index) for index in range(array.shape[spec.axis])]
+    vmin = float(np.min(array))
+    vmax = float(np.max(array))
+    figure, axis = plt.subplots(figsize=(6.4, 5.0), constrained_layout=True)
+    image = axis.imshow(frames[0], origin="lower", aspect="auto", cmap="viridis", vmin=vmin, vmax=vmax)
+    title = axis.set_title("")
+    axis.set_xlabel("poloidal index")
+    axis.set_ylabel("radial index")
+    figure.colorbar(image, ax=axis, shrink=0.88, pad=0.02, label=_display_label(field_name))
+
+    def _update(frame_index: int):
+        image.set_data(frames[frame_index])
+        coord = float(np.asarray(spec.coordinate_values, dtype=np.float64)[frame_index])
+        title.set_text(f"{_display_label(field_name)} · {_display_label(spec.name)} · {spec.coordinate_name}={coord:.3f}")
+        return (image, title)
+
+    anim = animation.FuncAnimation(figure, _update, frames=len(frames), interval=max(1, int(1000 / fps)), blit=False)
+    writer = animation.PillowWriter(fps=fps)
+    anim.save(target, writer=writer)
+    plt.close(figure)
+    return target
+
+
+def write_slice_report_json(report: dict[str, object], path: str | Path) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
+    return target
+
+
+def _extract_plane(values: np.ndarray, axis: int, index: int) -> np.ndarray:
+    if axis == 0:
+        return np.asarray(values[index, :, :], dtype=np.float64)
+    if axis == 1:
+        return np.asarray(values[:, index, :], dtype=np.float64)
+    if axis == 2:
+        return np.asarray(values[:, :, index], dtype=np.float64)
+    raise ValueError(f"Unsupported slice axis {axis}")
+
+
+def _display_label(name: str) -> str:
+    labels = {
+        "Bxy": "B",
+        "Bmag": "|B|",
+        "J": "Jacobian",
+        "jacobian": "Jacobian",
+        "g11": "g11",
+        "g22": "g22",
+        "g33": "g33",
+        "g_11": "g^11",
+        "g_22": "g^22",
+        "g_33": "g^33",
+        "radial_index_planes": "Radial slices",
+        "toroidal_index_planes": "Toroidal slices",
+        "poloidal_index_planes": "Poloidal slices",
+    }
+    return labels.get(name, name.replace("_", " "))

drbx/validation/publication_plotting.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from matplotlib import pyplot as plt
+import numpy as np
+
+
+PUBLICATION_DPI = 320
+
+
+def style_axis(
+    axis,
+    *,
+    title: str,
+    xlabel: str | None = None,
+    ylabel: str | None = None,
+    yscale: str | None = None,
+    xscale: str | None = None,
+    grid: str = "y",
+) -> None:
+    axis.set_title(title, fontsize=13.5, fontweight="semibold", pad=8.0)
+    if xlabel is not None:
+        axis.set_xlabel(xlabel, fontsize=12.0)
+    if ylabel is not None:
+        axis.set_ylabel(ylabel, fontsize=12.0)
+    if yscale is not None:
+        axis.set_yscale(yscale)
+    if xscale is not None:
+        axis.set_xscale(xscale)
+    axis.grid(alpha=0.22, axis=grid, linewidth=0.8)
+    axis.tick_params(labelsize=11.0)
+    axis.spines["top"].set_visible(False)
+    axis.spines["right"].set_visible(False)
+
+
+def annotate_bars(
+    axis,
+    x: np.ndarray,
+    values: np.ndarray,
+    *,
+    fmt: str = "{:.2e}",
+    fontsize: float = 9.0,
+    rotation: float = 0.0,
+) -> None:
+    values = np.asarray(values, dtype=np.float64)
+    if values.size == 0:
+        return
+    positive = np.abs(values[np.nonzero(values)])
+    scale = float(np.max(positive)) if positive.size else 1.0
+    offset = 0.03 * scale
+    for xi, value in zip(np.asarray(x, dtype=np.float64), values, strict=True):
+        va = "bottom" if value >= 0.0 else "top"
+        delta = offset if value >= 0.0 else -offset
+        axis.text(
+            float(xi),
+            float(value + delta),
+            fmt.format(float(value)),
+            ha="center",
+            va=va,
+            fontsize=fontsize,
+            rotation=rotation,
+        )
+
+
+def support_window_slice(
+    *arrays: np.ndarray,
+    padding: int = 4,
+    threshold: float = 1.0e-12,
+) -> slice:
+    support = np.zeros_like(np.asarray(arrays[0], dtype=np.float64), dtype=bool)
+    for array in arrays:
+        support |= np.abs(np.asarray(array, dtype=np.float64)) > threshold
+    indices = np.flatnonzero(support)
+    if indices.size == 0:
+        return slice(0, support.size)
+    start = max(0, int(indices[0]) - padding)
+    stop = min(support.size, int(indices[-1]) + padding + 1)
+    return slice(start, stop)
+
+
+def save_publication_figure(figure, output_path: str | Path) -> None:
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    figure.savefig(
+        output,
+        dpi=PUBLICATION_DPI,
+        bbox_inches="tight",
+        facecolor="white",
+    )
+    plt.close(figure)