11 — Neural-Seeded Predictive Equilibrium (Validated Warm Start)

This tutorial walks through the compiled predictive free-boundary equilibrium solver and the DeepONet predictive warm-starter that seeds it with a validated neural prediction.

The free-boundary Grad–Shafranov equilibrium is the fixed point of

$$\Delta^{*}\psi \;=\; -\mu_0 R^{2}\,p'(\psi)\;-\;F(\psi)F'(\psi),$$

with the plasma boundary determined self-consistently by the external poloidal-field coils. The predictive solver takes the coil currents together with the $p'(\psi)$ and $FF'(\psi)$ profile knots and iterates to the equilibrium that holds a target plasma current $I_p$.

The warm-starter does not replace this mechanistic solve. It provides an initial guess $\psi_{\text{init}}$ from a trained DeepONet, validates it against the mechanistic fixed point, and — if the neural inference is unstable or out of contract — falls back to the real cold solver. The neural component therefore can never degrade correctness: it is an accelerator candidate that is always checked, never trusted blindly.

License: © 1996–2026 Miroslav Šotek. GNU AGPL-3.0-or-later (commercial licence available).

Evidence boundary. This notebook is tutorial and onboarding material, not production-parity evidence. For accepted parity claims use the tracked reports under validation/reports/.

0 · Environment

The cell below installs the package when the notebook runs on a fresh Colab/Binder kernel and enables 64-bit JAX (required for the equilibrium fixed point).

The DeepONet warm-starter ships after the 4.0.0 release. When the installed wheel predates it, Sections 2–3 report that cleanly and are skipped; Section 1 (the compiled predictive solver, part of 4.0.0) runs everywhere.

# --- Colab/Binder bootstrap: no-op in a repo checkout or CI ---
try:
    import scpn_fusion  # noqa: F401
except ImportError:
    import subprocess, sys
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "scpn-fusion"], check=True)

import importlib.util
from typing import Any, cast
import jax
cast(Any, jax.config).update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np

# The warm-starter is post-4.0.0; degrade gracefully on an older wheel.
HAS_WARM_START = importlib.util.find_spec(
    "scpn_fusion.core.deeponet_solver_warm_start"
) is not None
print("scpn_fusion warm-start API available:", HAS_WARM_START)

scpn_fusion warm-start API available: True

1 · Compiled predictive free-boundary solve

We use the synthetic diverted case from the predictive test suite verbatim: six poloidal-field coils on a $33\times33$ $(R,Z)$ grid, five-knot $p'$ and $FF'$ profiles, and a $1\,\mathrm{MA}$ target. solve_predictive_equilibrium_compiled returns the equilibrium flux $\psi(R,Z)$ and, with return_iterations=True, the number of fixed-point iterations to convergence.

from scpn_fusion.core.jax_free_boundary_predictive import build_response_matrix
from scpn_fusion.core.jax_predictive_forward_compiled import (
    solve_predictive_equilibrium_compiled,
)

R = jnp.linspace(1.0, 2.5, 33)
Z = jnp.linspace(-1.4, 1.4, 33)
COIL_R = jnp.asarray([1.2, 2.3, 1.2, 2.3, 1.6, 1.5])
COIL_Z = jnp.asarray([0.9, 0.9, -0.9, -0.9, 1.3, -1.35])
COIL_I = jnp.asarray([-3.0e5, -3.0e5, -3.0e5, -3.0e5, -1.0e5, -6.0e5])
COIL_NAMES = ("C1", "C2", "C3", "C4", "C5", "C6")
PSIN = jnp.linspace(0.0, 1.0, 6)
PPRIME = jnp.asarray([-8.0e4, -6.0e4, -4.0e4, -2.0e4, -0.7e4, 0.0])
FFPRIME = jnp.asarray([-1.2, -0.9, -0.6, -0.3, -0.1, 0.0])
IP_TARGET = 1.0e6

response_matrix, wall_idx, source_idx = build_response_matrix(R, Z)

psi_cold, iters_cold = cast(
    tuple,
    solve_predictive_equilibrium_compiled(
        COIL_I, PPRIME, FFPRIME, R, Z, COIL_R, COIL_Z, PSIN, IP_TARGET,
        response_matrix, wall_idx, source_idx,
        n_iter=150, return_iterations=True,
    ),
)
psi_cold.block_until_ready()
flux_span = float(jnp.ptp(psi_cold))
print(f"cold predictive solve: {int(iters_cold)} iterations, flux span = {flux_span:.4g} Wb/rad")

WARNING:2026-08-29 20:27:31,340:jax._src.xla_bridge:864: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

cold predictive solve: 131 iterations, flux span = 0.7617 Wb/rad

import matplotlib.pyplot as plt

psi_np = np.asarray(psi_cold)
Rg, Zg = np.meshgrid(np.asarray(R), np.asarray(Z), indexing="xy")
fig, ax = plt.subplots(figsize=(4.6, 5.2))
cf = ax.contourf(Rg, Zg, psi_np, levels=40, cmap="viridis")
ax.contour(Rg, Zg, psi_np, levels=14, colors="white", linewidths=0.4, alpha=0.6)
ax.scatter(np.asarray(COIL_R), np.asarray(COIL_Z), c="crimson", marker="s", s=40,
           label="PF coils", zorder=5)
ax.set_aspect("equal")
ax.set_xlabel("R  [m]"); ax.set_ylabel("Z  [m]")
ax.set_title("Predictive free-boundary equilibrium  ψ(R, Z)")
ax.legend(loc="upper right", fontsize=8)
fig.colorbar(cf, ax=ax, label="ψ  [Wb/rad]", shrink=0.85)
plt.tight_layout(); plt.show()

2 · Validated neural warm start

The warm-starter binds one authenticated DeepONet artifact to a machine manifest, predicts a seed field, and hands it to the same compiled solver as psi_init. Here we build a self-contained artifact whose predicted field is the equilibrium of Section 1 — an exact seed — so we can show that the neural path reaches the same mechanistic fixed point (relative difference below the solver tolerance). No external weights are required.

if not HAS_WARM_START:
    print("Warm-start API not in this wheel (pre-release). "
          "Install scpn-fusion from a checkout of main to run Sections 2-3.")
else:
    import tempfile
    from pathlib import Path
    from scpn_fusion.core import (
        DeepONetPredictiveWarmStarter, DeepONetWarmStartResult,
    )

    SHA = "b" * 64
    FEATURES = (
        "plasma_current_target_a",
        *(f"coil_current_a.{n}" for n in COIL_NAMES),
        *(f"pprime_knot_{i}" for i in range(5)),
        *(f"ffprime_knot_{i}" for i in range(5)),
    )

    def write_seed_artifact(path, field_mean, *, unstable=False):
        gr, gz = np.meshgrid(np.asarray(R), np.asarray(Z), indexing="xy")
        coords = np.column_stack((gr.ravel(), gz.ravel()))
        branch_w = np.zeros((len(FEATURES), 1))
        if unstable:
            branch_w.fill(np.finfo(np.float64).max)
        np.savez(
            path,
            artifact_schema=np.asarray(["scpn-fusion.equilibrium-deeponet.v1"]),
            branch_n_layers=np.asarray([1]), branch_0_W=branch_w, branch_0_b=np.zeros(1),
            trunk_n_layers=np.asarray([1]), trunk_0_W=np.ones((2, 1)), trunk_0_b=np.zeros(1),
            input_mean=np.zeros(len(FEATURES)), input_std=np.ones(len(FEATURES)),
            coordinates_rz_m=coords, coordinate_mean=np.mean(coords, 0),
            coordinate_std=np.std(coords, 0),
            field_mean=np.asarray(field_mean).ravel(), field_scale=np.asarray([1.0]),
            basis_width=np.asarray([1]), grid_nh=np.asarray([len(Z)]),
            grid_nw=np.asarray([len(R)]), feature_names=np.asarray(FEATURES),
            dataset_manifest_sha256=np.asarray([SHA]),
        )

    solve_kwargs = dict(
        machine_manifest_sha256=SHA, coil_names=COIL_NAMES, coil_i=COIL_I,
        pprime_vals=PPRIME, ffprime_vals=FFPRIME, r_grid=R, z_grid=Z,
        coil_r=COIL_R, coil_z=COIL_Z, psin_knots=PSIN, ip_target=IP_TARGET,
        response_matrix=response_matrix, wall_idx=wall_idx, source_idx=source_idx,
        n_iter=150,
    )

    tmp = Path(tempfile.mkdtemp())
    write_seed_artifact(tmp / "exact_seed.npz", psi_cold)
    starter = DeepONetPredictiveWarmStarter(tmp / "exact_seed.npz", prefer_rust=False)
    result = cast(DeepONetWarmStartResult, starter.solve(**solve_kwargs))
    result.equilibrium.block_until_ready()

    rel = float(jnp.max(jnp.abs(result.equilibrium - psi_cold))) / flux_span
    print(f"used_neural_seed   : {result.used_neural_seed}")
    print(f"fallback_reason    : {result.fallback_reason}")
    print(f"iterations         : {result.iterations}")
    print(f"neural_seed_weight : {result.neural_seed_weight}")
    print(f"rel. diff to cold  : {rel:.2e}  (< solver tolerance)")

used_neural_seed   : True
fallback_reason    : None
iterations         : 131
neural_seed_weight : 1.0
rel. diff to cold  : 4.37e-16  (< solver tolerance)

3 · Guaranteed mechanistic fallback

Safety is the point of the design. We now hand the warm-starter a corrupt artifact whose branch weights are non-finite. The neural inference fails its finiteness contract, the warm-starter reports used_neural_seed = False with a machine-readable reason, and still returns the correct mechanistic equilibrium by running the cold solver. A broken or adversarial surrogate cannot make the answer worse than the mechanistic baseline.

if HAS_WARM_START:
    write_seed_artifact(tmp / "bad_seed.npz", psi_cold, unstable=True)
    bad_starter = DeepONetPredictiveWarmStarter(tmp / "bad_seed.npz", prefer_rust=False)
    fallback = cast(DeepONetWarmStartResult, bad_starter.solve(**solve_kwargs))
    fallback.equilibrium.block_until_ready()

    rel_fb = float(jnp.max(jnp.abs(fallback.equilibrium - psi_cold))) / flux_span
    print(f"used_neural_seed : {fallback.used_neural_seed}")
    print(f"fallback_reason  : {fallback.fallback_reason!r}")
    print(f"rel. diff to cold: {rel_fb:.2e}  (mechanistic equilibrium preserved)")
else:
    print("Skipped: warm-start API not available in this wheel.")

used_neural_seed : False
fallback_reason  : 'neural_inference_failed'
rel. diff to cold: 0.00e+00  (mechanistic equilibrium preserved)

What this shows — and what it does not

• Shown. The compiled predictive solver reaches the free-boundary Grad–Shafranov fixed point that holds the target $I_p$; a validated DeepONet seed reaches the same fixed point to well below solver tolerance; and a failed/adversarial seed transparently falls back to the mechanistic solve.
• Not claimed here. This tutorial makes no runtime-speedup claim. The synthetic case is dominated by the separatrix-continuation schedule, so iteration counts are comparable by design; the warm-starter's contribution demonstrated here is validated seeding with a guaranteed fallback, not acceleration. Accepted performance and parity claims live in the tracked reports under validation/reports/.

See also notebook 03_grad_shafranov_equilibrium for the eager solver contract and 06_inverse_and_transport_benchmarks for the inverse/transport validation.