drbx 2.0.0.dev0__py3-none-any.whl

drbx/validation/stellarator_vorticity_campaign.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from matplotlib import pyplot as plt
+import jax.numpy as jnp
+import numpy as np
+
+from ..geometry import build_synthetic_stellarator_geometry
+from ..native.fci_vorticity import apply_fci_vorticity_operator, solve_fci_vorticity_potential_cg
+
+
+@dataclass(frozen=True)
+class StellaratorVorticityCampaignArtifacts:
+    report_json_path: Path
+    arrays_npz_path: Path
+    plot_png_path: Path
+
+
+def create_stellarator_vorticity_campaign_package(
+    *,
+    output_root: str | Path,
+    case_label: str = "stellarator_vorticity_campaign",
+    nx: int = 28,
+    ny: int = 26,
+    nz: int = 52,
+) -> StellaratorVorticityCampaignArtifacts:
+    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, arrays = build_stellarator_vorticity_campaign(nx=nx, ny=ny, nz=nz)
+    report_json_path = data_dir / f"{case_label}.json"
+    report_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, **arrays)
+    plot_png_path = images_dir / f"{case_label}.png"
+    save_stellarator_vorticity_plot(report, arrays, plot_png_path)
+    return StellaratorVorticityCampaignArtifacts(
+        report_json_path=report_json_path,
+        arrays_npz_path=arrays_npz_path,
+        plot_png_path=plot_png_path,
+    )
+
+
+def build_stellarator_vorticity_campaign(
+    *,
+    nx: int = 28,
+    ny: int = 26,
+    nz: int = 52,
+) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
+    geometry = build_synthetic_stellarator_geometry(nx=nx, ny=ny, nz=nz)
+    phi_exact = _manufactured_potential(geometry)
+    density = _manufactured_density(geometry)
+    boussinesq_vorticity = apply_fci_vorticity_operator(
+        phi_exact,
+        density,
+        geometry.metric,
+        boussinesq=True,
+    )
+    non_boussinesq_vorticity = apply_fci_vorticity_operator(
+        phi_exact,
+        density,
+        geometry.metric,
+        boussinesq=False,
+    )
+    boussinesq_solve = solve_fci_vorticity_potential_cg(
+        boussinesq_vorticity,
+        density,
+        geometry.metric,
+        iterations=600,
+        boussinesq=True,
+    )
+    non_boussinesq_solve = solve_fci_vorticity_potential_cg(
+        non_boussinesq_vorticity,
+        density,
+        geometry.metric,
+        iterations=600,
+        boussinesq=False,
+    )
+    phi_centered = _remove_weighted_mean(np.asarray(phi_exact), np.asarray(geometry.metric.J))
+    boussinesq_phi = np.asarray(boussinesq_solve.potential, dtype=np.float64)
+    non_boussinesq_phi = np.asarray(non_boussinesq_solve.potential, dtype=np.float64)
+    boussinesq_error = boussinesq_phi - phi_centered
+    non_boussinesq_error = non_boussinesq_phi - phi_centered
+    phi_norm = max(np.sqrt(np.mean(phi_centered * phi_centered)), 1.0e-30)
+    boussinesq_relative_l2_error = float(
+        np.sqrt(np.mean(boussinesq_error * boussinesq_error)) / phi_norm
+    )
+    non_boussinesq_relative_l2_error = float(
+        np.sqrt(np.mean(non_boussinesq_error * non_boussinesq_error)) / phi_norm
+    )
+    boussinesq_residual_l2 = float(np.asarray(boussinesq_solve.residual_l2))
+    non_boussinesq_residual_l2 = float(np.asarray(non_boussinesq_solve.residual_l2))
+    boussinesq_vorticity_np = np.asarray(boussinesq_vorticity, dtype=np.float64)
+    non_boussinesq_vorticity_np = np.asarray(non_boussinesq_vorticity, dtype=np.float64)
+    vorticity_difference = non_boussinesq_vorticity_np - boussinesq_vorticity_np
+    operator_difference_relative_l2 = float(
+        np.sqrt(np.mean(vorticity_difference * vorticity_difference))
+        / max(np.sqrt(np.mean(boussinesq_vorticity_np * boussinesq_vorticity_np)), 1.0e-30)
+    )
+    coefficient = np.asarray(density / jnp.maximum(jnp.square(geometry.metric.Bxy), 1.0e-30))
+    coefficient_contrast = float(np.max(coefficient) / max(float(np.min(coefficient)), 1.0e-30))
+    constant_coefficient_density = 1.7 * jnp.square(geometry.metric.Bxy)
+    constant_boussinesq_operator = apply_fci_vorticity_operator(
+        phi_exact,
+        constant_coefficient_density,
+        geometry.metric,
+        boussinesq=True,
+    )
+    constant_non_boussinesq_operator = apply_fci_vorticity_operator(
+        phi_exact,
+        constant_coefficient_density,
+        geometry.metric,
+        boussinesq=False,
+    )
+    constant_coefficient_operator_linf = float(
+        np.max(
+            np.abs(
+                np.asarray(constant_boussinesq_operator)
+                - np.asarray(constant_non_boussinesq_operator)
+            )
+        )
+    )
+    exb_radial = _radial_exb_proxy(non_boussinesq_phi, geometry)
+    report: dict[str, Any] = {
+        "case": "non_axisymmetric_boussinesq_non_boussinesq_vorticity_inversion",
+        "geometry": geometry.metadata,
+        "iterations": int(non_boussinesq_solve.iterations),
+        "boussinesq_relative_l2_potential_error": boussinesq_relative_l2_error,
+        "non_boussinesq_relative_l2_potential_error": non_boussinesq_relative_l2_error,
+        "relative_l2_potential_error": non_boussinesq_relative_l2_error,
+        "boussinesq_relative_residual_l2": boussinesq_residual_l2,
+        "non_boussinesq_relative_residual_l2": non_boussinesq_residual_l2,
+        "relative_residual_l2": non_boussinesq_residual_l2,
+        "operator_difference_relative_l2": operator_difference_relative_l2,
+        "constant_coefficient_operator_linf": constant_coefficient_operator_linf,
+        "density_min": float(np.min(np.asarray(density))),
+        "density_max": float(np.max(np.asarray(density))),
+        "density_over_b_squared_contrast": coefficient_contrast,
+        "radial_exb_proxy_rms": float(np.sqrt(np.mean(exb_radial * exb_radial))),
+        "potential_rms": float(np.sqrt(np.mean(non_boussinesq_phi * non_boussinesq_phi))),
+        "boussinesq_vorticity_rms": float(
+            np.sqrt(np.mean(boussinesq_vorticity_np * boussinesq_vorticity_np))
+        ),
+        "non_boussinesq_vorticity_rms": float(
+            np.sqrt(np.mean(non_boussinesq_vorticity_np * non_boussinesq_vorticity_np))
+        ),
+        "vorticity_rms": float(
+            np.sqrt(np.mean(non_boussinesq_vorticity_np * non_boussinesq_vorticity_np))
+        ),
+    }
+    report["passed"] = (
+        boussinesq_relative_l2_error < 2.5e-2
+        and non_boussinesq_relative_l2_error < 2.5e-2
+        and boussinesq_residual_l2 < 5.0e-3
+        and non_boussinesq_residual_l2 < 5.0e-3
+        and operator_difference_relative_l2 > 5.0e-2
+        and constant_coefficient_operator_linf < 1.0e-8
+        and report["radial_exb_proxy_rms"] > 1.0e-4
+    )
+    arrays = {
+        "phi_exact_slice": phi_centered[:, 0, :].astype(np.float32),
+        "phi_solved_slice": non_boussinesq_phi[:, 0, :].astype(np.float32),
+        "boussinesq_phi_solved_slice": boussinesq_phi[:, 0, :].astype(np.float32),
+        "non_boussinesq_phi_solved_slice": non_boussinesq_phi[:, 0, :].astype(np.float32),
+        "phi_error_slice": non_boussinesq_error[:, 0, :].astype(np.float32),
+        "boussinesq_phi_error_slice": boussinesq_error[:, 0, :].astype(np.float32),
+        "non_boussinesq_phi_error_slice": non_boussinesq_error[:, 0, :].astype(np.float32),
+        "vorticity_slice": non_boussinesq_vorticity_np[:, 0, :].astype(np.float32),
+        "boussinesq_vorticity_slice": boussinesq_vorticity_np[:, 0, :].astype(np.float32),
+        "non_boussinesq_vorticity_slice": non_boussinesq_vorticity_np[:, 0, :].astype(np.float32),
+        "vorticity_difference_slice": vorticity_difference[:, 0, :].astype(np.float32),
+        "density_over_b_squared_slice": coefficient[:, 0, :].astype(np.float32),
+        "radial_exb_slice": exb_radial[:, 0, :].astype(np.float32),
+        "summary": np.asarray(
+            [
+                boussinesq_relative_l2_error,
+                non_boussinesq_relative_l2_error,
+                boussinesq_residual_l2,
+                non_boussinesq_residual_l2,
+                operator_difference_relative_l2,
+                constant_coefficient_operator_linf,
+                report["radial_exb_proxy_rms"],
+                report["potential_rms"],
+            ],
+            dtype=np.float32,
+        ),
+    }
+    return report, arrays
+
+
+def save_stellarator_vorticity_plot(
+    report: dict[str, Any],
+    arrays: dict[str, np.ndarray],
+    path: str | Path,
+) -> Path:
+    resolved = Path(path)
+    resolved.parent.mkdir(parents=True, exist_ok=True)
+    fig, axes = plt.subplots(3, 3, figsize=(16.2, 12.2), constrained_layout=True)
+    image0 = axes[0, 0].imshow(arrays["phi_exact_slice"], origin="lower", aspect="auto", cmap="viridis")
+    axes[0, 0].set_title("manufactured potential")
+    fig.colorbar(image0, ax=axes[0, 0])
+    image1 = axes[0, 1].imshow(arrays["boussinesq_phi_solved_slice"], origin="lower", aspect="auto", cmap="viridis")
+    axes[0, 1].set_title("Boussinesq reconstructed potential")
+    fig.colorbar(image1, ax=axes[0, 1])
+    image2 = axes[0, 2].imshow(arrays["non_boussinesq_phi_solved_slice"], origin="lower", aspect="auto", cmap="viridis")
+    axes[0, 2].set_title("non-Boussinesq reconstructed potential")
+    fig.colorbar(image2, ax=axes[0, 2])
+    image3 = axes[1, 0].imshow(arrays["boussinesq_vorticity_slice"], origin="lower", aspect="auto", cmap="magma")
+    axes[1, 0].set_title("Boussinesq vorticity")
+    fig.colorbar(image3, ax=axes[1, 0])
+    image4 = axes[1, 1].imshow(arrays["non_boussinesq_vorticity_slice"], origin="lower", aspect="auto", cmap="magma")
+    axes[1, 1].set_title("non-Boussinesq vorticity")
+    fig.colorbar(image4, ax=axes[1, 1])
+    vmax_vort = float(np.max(np.abs(arrays["vorticity_difference_slice"])))
+    image5 = axes[1, 2].imshow(
+        arrays["vorticity_difference_slice"],
+        origin="lower",
+        aspect="auto",
+        cmap="coolwarm",
+        vmin=-vmax_vort,
+        vmax=vmax_vort,
+    )
+    axes[1, 2].set_title("non-Boussinesq minus Boussinesq")
+    fig.colorbar(image5, ax=axes[1, 2])
+    vmax_bouss = float(np.max(np.abs(arrays["boussinesq_phi_error_slice"])))
+    image6 = axes[2, 0].imshow(
+        arrays["boussinesq_phi_error_slice"],
+        origin="lower",
+        aspect="auto",
+        cmap="coolwarm",
+        vmin=-vmax_bouss,
+        vmax=vmax_bouss,
+    )
+    axes[2, 0].set_title("Boussinesq potential error")
+    fig.colorbar(image6, ax=axes[2, 0])
+    vmax_non_bouss = float(np.max(np.abs(arrays["non_boussinesq_phi_error_slice"])))
+    image7 = axes[2, 1].imshow(
+        arrays["non_boussinesq_phi_error_slice"],
+        origin="lower",
+        aspect="auto",
+        cmap="coolwarm",
+        vmin=-vmax_non_bouss,
+        vmax=vmax_non_bouss,
+    )
+    axes[2, 1].set_title("non-Boussinesq potential error")
+    fig.colorbar(image7, ax=axes[2, 1])
+    labels = [
+        "Bq err",
+        "non-Bq err",
+        "Bq res",
+        "non-Bq res",
+        "op diff",
+        "const eq",
+        "ExB rms",
+        "phi rms",
+    ]
+    colors = ["#9b2226", "#bb3e03", "#0a9396", "#005f73", "#ee9b00", "#94d2bd", "#ca6702", "#001219"]
+    axes[2, 2].bar(np.arange(len(labels)), arrays["summary"], color=colors)
+    axes[2, 2].set_xticks(np.arange(len(labels)), labels, rotation=28, ha="right")
+    axes[2, 2].set_yscale("log")
+    axes[2, 2].grid(axis="y", alpha=0.25)
+    axes[2, 2].set_title("model and inversion metrics")
+    for axis in axes.ravel()[:8]:
+        axis.set_xlabel("poloidal index")
+        axis.set_ylabel("radial index")
+    fig.suptitle(
+        "Non-axisymmetric vorticity gate: Boussinesq vs non-Boussinesq perpendicular polarization",
+        fontsize=14,
+    )
+    fig.savefig(resolved, dpi=190)
+    plt.close(fig)
+    return resolved
+
+
+def _manufactured_potential(geometry: object) -> jnp.ndarray:
+    return (
+        jnp.sin(jnp.pi * geometry.radial) * jnp.cos(2.0 * geometry.poloidal_angle - 5.0 * geometry.toroidal_angle)
+        + 0.20 * geometry.radial * jnp.sin(3.0 * geometry.poloidal_angle + 2.0 * geometry.toroidal_angle)
+    )
+
+
+def _manufactured_density(geometry: object) -> jnp.ndarray:
+    density = (
+        1.0
+        + 0.35 * geometry.radial
+        + 0.12 * jnp.cos(geometry.poloidal_angle - 2.0 * geometry.toroidal_angle)
+        + 0.08 * jnp.cos(3.0 * geometry.toroidal_angle)
+    )
+    return jnp.maximum(density, 0.1)
+
+
+def _remove_weighted_mean(value: np.ndarray, weights: np.ndarray) -> np.ndarray:
+    return value - float(np.sum(weights * value) / np.sum(weights))
+
+
+def _radial_exb_proxy(phi: np.ndarray, geometry: object) -> np.ndarray:
+    dz = float(2.0 * np.pi / geometry.shape[2])
+    return -(np.roll(phi, -1, axis=2) - np.roll(phi, 1, axis=2)) / (2.0 * dz)

drbx/validation/vmec_extender_edge_field_campaign.py

@@ -0,0 +1,260 @@
+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 ..geometry.vmec_extender_import import (
+    VmecExtenderGrid,
+    interpolate_vmec_extender_B_cyl,
+    load_vmec_extender_grid_netcdf,
+    vmec_extender_fieldline_rhs_RZ_phi,
+)
+from .publication_plotting import annotate_bars, save_publication_figure, style_axis
+
+
+@dataclass(frozen=True)
+class VmecExtenderEdgeFieldCampaignArtifacts:
+    summary_json_path: Path
+    arrays_npz_path: Path
+    plot_png_path: Path
+
+
+def create_vmec_extender_edge_field_campaign_package(
+    *,
+    output_root: str | Path,
+    field_grid_path: str | Path,
+    case_label: str = "vmec_extender_edge_field_campaign",
+    strict_metadata: bool = True,
+) -> VmecExtenderEdgeFieldCampaignArtifacts:
+    """Run the imported-field verification campaign and write public artifacts."""
+
+    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_vmec_extender_edge_field_campaign_report(
+        field_grid_path=field_grid_path,
+        strict_metadata=strict_metadata,
+    )
+    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,
+        metric_names=np.asarray(
+            [
+                "node_interpolation_max_abs_error",
+                "midpoint_interpolation_max_abs_error",
+                "field_period_relative_l2",
+                "fieldline_rhs_max_abs_error",
+                "absB_consistency_max_abs_error",
+            ],
+            dtype=object,
+        ),
+        metric_values=np.asarray(
+            [
+                report["node_interpolation_max_abs_error"],
+                report["midpoint_interpolation_max_abs_error"],
+                report["field_period_relative_l2"],
+                report["fieldline_rhs_max_abs_error"],
+                report["absB_consistency_max_abs_error"],
+            ],
+            dtype=np.float64,
+        ),
+        node_points=np.asarray(report["node_sample_points"], dtype=np.float64),
+        node_expected_B=np.asarray(report["node_sample_expected_B"], dtype=np.float64),
+        node_interpolated_B=np.asarray(report["node_sample_interpolated_B"], dtype=np.float64),
+    )
+
+    plot_png_path = save_vmec_extender_edge_field_campaign_plot(report, images_dir / f"{case_label}.png")
+    return VmecExtenderEdgeFieldCampaignArtifacts(
+        summary_json_path=summary_json_path,
+        arrays_npz_path=arrays_npz_path,
+        plot_png_path=plot_png_path,
+    )
+
+
+def build_vmec_extender_edge_field_campaign_report(
+    *,
+    field_grid_path: str | Path,
+    strict_metadata: bool = True,
+) -> dict[str, object]:
+    """Build a numerical verification report for an imported edge-field grid."""
+
+    grid = load_vmec_extender_grid_netcdf(field_grid_path, strict_metadata=strict_metadata)
+    node_points = _node_points(grid)
+    node_expected = _field_stack(grid)
+    node_interpolated = np.asarray(interpolate_vmec_extender_B_cyl(grid, node_points), dtype=np.float64)
+    node_error = float(np.max(np.abs(node_interpolated - node_expected)))
+
+    midpoint_points, midpoint_expected = _midpoint_trilinear_reference(grid)
+    midpoint_interpolated = np.asarray(interpolate_vmec_extender_B_cyl(grid, midpoint_points), dtype=np.float64)
+    midpoint_error = float(np.max(np.abs(midpoint_interpolated - midpoint_expected))) if midpoint_expected.size else 0.0
+
+    period_points = _periodicity_probe_points(grid)
+    period_base = np.asarray(interpolate_vmec_extender_B_cyl(grid, period_points), dtype=np.float64)
+    shifted_points = period_points.copy()
+    shifted_points[..., 1] = shifted_points[..., 1] + float(grid.phi_period)
+    period_shifted = np.asarray(interpolate_vmec_extender_B_cyl(grid, shifted_points), dtype=np.float64)
+    period_relative_l2 = _relative_l2(period_base, period_shifted)
+
+    rhs_points = midpoint_points.reshape((-1, 3))[: max(1, min(32, midpoint_points.reshape((-1, 3)).shape[0]))]
+    rhs_B = np.asarray(interpolate_vmec_extender_B_cyl(grid, rhs_points), dtype=np.float64)
+    rhs_expected = np.stack(
+        (
+            rhs_points[:, 0] * rhs_B[:, 0] / rhs_B[:, 1],
+            rhs_points[:, 0] * rhs_B[:, 2] / rhs_B[:, 1],
+        ),
+        axis=-1,
+    )
+    rhs_actual = np.asarray(vmec_extender_fieldline_rhs_RZ_phi(grid, rhs_points), dtype=np.float64)
+    rhs_error = float(np.max(np.abs(rhs_actual - rhs_expected)))
+
+    absB_error = float(grid.metadata.get("absB_consistency_max_abs_error", np.nan))
+    passed = (
+        node_error <= 1.0e-10
+        and midpoint_error <= 1.0e-10
+        and period_relative_l2 <= 1.0e-12
+        and rhs_error <= 1.0e-10
+        and absB_error <= 1.0e-8
+    )
+    sample_count = min(8, node_points.reshape((-1, 3)).shape[0])
+    return {
+        "family": "vmec_extender_edge_field_campaign",
+        "case": Path(field_grid_path).stem,
+        "source": str(grid.metadata.get("source", "unknown")),
+        "grid_shape": [int(value) for value in grid.shape],
+        "nfp": int(grid.nfp),
+        "phi_period": float(grid.phi_period),
+        "metadata_passed": True,
+        "node_interpolation_max_abs_error": node_error,
+        "midpoint_interpolation_max_abs_error": midpoint_error,
+        "field_period_relative_l2": period_relative_l2,
+        "fieldline_rhs_max_abs_error": rhs_error,
+        "absB_consistency_max_abs_error": absB_error,
+        "passed": bool(passed),
+        "notes": (
+            "The campaign verifies the gridded VMEC-extender import contract, "
+            "physical-phi wrapping, trilinear interpolation, and field-line RHS."
+        ),
+        "node_sample_points": node_points.reshape((-1, 3))[:sample_count].tolist(),
+        "node_sample_expected_B": node_expected.reshape((-1, 3))[:sample_count].tolist(),
+        "node_sample_interpolated_B": node_interpolated.reshape((-1, 3))[:sample_count].tolist(),
+    }
+
+
+def save_vmec_extender_edge_field_campaign_plot(report: dict[str, object], path: str | Path) -> Path:
+    """Save a compact publication-style summary plot for the import campaign."""
+
+    metrics = np.asarray(
+        [
+            report["node_interpolation_max_abs_error"],
+            report["midpoint_interpolation_max_abs_error"],
+            report["field_period_relative_l2"],
+            report["fieldline_rhs_max_abs_error"],
+            report["absB_consistency_max_abs_error"],
+        ],
+        dtype=np.float64,
+    )
+    labels = [
+        "node\ninterp",
+        "midpoint\ninterp",
+        "field-period\nwrap",
+        "field-line\nRHS",
+        "|B|\nclosure",
+    ]
+    floor = np.maximum(metrics, 1.0e-16)
+    colors = ["#0a9396" if bool(report["passed"]) else "#bb3e03"] * len(metrics)
+    figure, axis = plt.subplots(figsize=(7.2, 3.8), constrained_layout=True)
+    x = np.arange(len(metrics), dtype=np.float64)
+    axis.bar(x, floor, color=colors, width=0.66, alpha=0.92)
+    axis.set_xticks(x, labels)
+    style_axis(
+        axis,
+        title="VMEC-extender edge-field import",
+        ylabel="error metric",
+        yscale="log",
+    )
+    axis.tick_params(axis="x", labelsize=8.5)
+    axis.tick_params(axis="y", labelsize=8.5)
+    axis.set_title("VMEC-extender edge-field import", fontsize=11.0, fontweight="semibold", pad=7.0)
+    axis.set_ylabel("error metric", fontsize=9.5)
+    annotate_bars(axis, x, floor, fmt="{:.1e}", fontsize=7.5)
+    axis.text(
+        0.98,
+        0.95,
+        f"shape={tuple(report['grid_shape'])}\nnfp={report['nfp']}, period={float(report['phi_period']):.4f}",
+        transform=axis.transAxes,
+        va="top",
+        ha="right",
+        fontsize=8.0,
+        linespacing=1.25,
+        bbox={"facecolor": "white", "alpha": 0.9, "edgecolor": "#d9d9d9", "linewidth": 0.6},
+    )
+    save_publication_figure(figure, path)
+    return Path(path)
+
+
+def _node_points(grid: VmecExtenderGrid) -> np.ndarray:
+    R, phi, Z = np.meshgrid(np.asarray(grid.R), np.asarray(grid.phi), np.asarray(grid.Z), indexing="ij")
+    return np.stack((R, phi, Z), axis=-1)
+
+
+def _field_stack(grid: VmecExtenderGrid) -> np.ndarray:
+    return np.stack(
+        (
+            np.asarray(grid.BR, dtype=np.float64),
+            np.asarray(grid.Bphi, dtype=np.float64),
+            np.asarray(grid.BZ, dtype=np.float64),
+        ),
+        axis=-1,
+    )
+
+
+def _midpoint_trilinear_reference(grid: VmecExtenderGrid) -> tuple[np.ndarray, np.ndarray]:
+    R = np.asarray(grid.R, dtype=np.float64)
+    phi = np.asarray(grid.phi, dtype=np.float64)
+    Z = np.asarray(grid.Z, dtype=np.float64)
+    if min(R.size, phi.size, Z.size) < 2:
+        return np.zeros((0, 3), dtype=np.float64), np.zeros((0, 3), dtype=np.float64)
+    R_mid = 0.5 * (R[:-1] + R[1:])
+    phi_mid = 0.5 * (phi[:-1] + phi[1:])
+    Z_mid = 0.5 * (Z[:-1] + Z[1:])
+    RR, PP, ZZ = np.meshgrid(R_mid, phi_mid, Z_mid, indexing="ij")
+    points = np.stack((RR, PP, ZZ), axis=-1)
+    fields = _field_stack(grid)
+    expected = 0.125 * (
+        fields[:-1, :-1, :-1, :]
+        + fields[1:, :-1, :-1, :]
+        + fields[:-1, 1:, :-1, :]
+        + fields[1:, 1:, :-1, :]
+        + fields[:-1, :-1, 1:, :]
+        + fields[1:, :-1, 1:, :]
+        + fields[:-1, 1:, 1:, :]
+        + fields[1:, 1:, 1:, :]
+    )
+    return points, expected
+
+
+def _periodicity_probe_points(grid: VmecExtenderGrid) -> np.ndarray:
+    R = np.asarray(grid.R, dtype=np.float64)
+    phi = np.asarray(grid.phi, dtype=np.float64)
+    Z = np.asarray(grid.Z, dtype=np.float64)
+    R_probe = 0.5 * (R[:-1] + R[1:])
+    phi_probe = phi[: max(1, min(3, phi.size))]
+    Z_probe = 0.5 * (Z[:-1] + Z[1:])
+    RR, PP, ZZ = np.meshgrid(R_probe, phi_probe, Z_probe, indexing="ij")
+    return np.stack((RR, PP, ZZ), axis=-1)
+
+
+def _relative_l2(expected: np.ndarray, actual: np.ndarray) -> float:
+    numerator = float(np.linalg.norm(np.ravel(actual - expected)))
+    denominator = float(max(np.linalg.norm(np.ravel(expected)), np.finfo(np.float64).tiny))
+    return numerator / denominator