examples: adjoint boundary sensitivities + SIMSOPT analytic gradient

examples/data/README.md

@@ -1,13 +1,13 @@
 # Example input and output files for VMEC++
 
-A few cases are available for testing VMCE++ and experimenting with it:
+A few cases are available for testing VMEC++ and experimenting with it:
 
 1. `cth_like_fixed_bdy.json` - Stellarator case, similar to the Compact Toroidal Hybrid ([CTH](https://www.auburn.edu/cosam/departments/physics/research/plasma_physics/compact_toroidal_hybrid/index.htm)) device
    1. `input.cth_like_fixed_bdy` - Fortran namelist input file to be used with Fortran VMEC
    1. `cth_like_fixed_bdy.json` - JSON input file for VMEC++, derived from `input.cth_like_fixed_bdy` using [`indata2json`](https://github.com/jonathanschilling/indata2json)
    1. `wout_cth_like_fixed_bdy.nc` - NetCDF output file, produced using [`PARVMEC`](https://github.com/ORNL-Fusion/PARVMEC) from `input.cth_like_fixed_bdy`, for testing the loading of a `wout` file using VMEC++'s tooling
 
-1. `input.nfp4_QH_warm_start` - quasi-helically example for use with SIMSOPT
+1. `input.nfp4_QH_warm_start` - quasi-helical example for use with SIMSOPT
 
 1. `solovev` - axisymmetric Tokamak case, similar to the Solov'ev equilibrium used in the [1983 Hirshman & Whitson article](https://doi.org/10.1063/1.864116)
     1. `input.solovev` - Fortran namelist input file for use with Fortran VMEC

examples/external_optimizers.py

@@ -0,0 +1,283 @@
+# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH
+# <info@proximafusion.com>
+#
+# SPDX-License-Identifier: MIT
+"""Drive a VMEC++ equilibrium from outside with general-purpose optimizers.
+
+VMEC's equilibrium is the stationary point of its augmented functional (MHD
+energy plus the spectral-condensation and lambda constraints). The gradient of
+that functional in the decomposed internal basis is the raw, unpreconditioned
+force exposed by ``VmecModel.evaluate(precondition=False)``. Finding the
+equilibrium is therefore the root problem F(x) = 0, which gradient- and
+Hessian-based solvers can attack while reusing VMEC++'s forward model, its
+preconditioner (``apply_preconditioner``, VMEC's approximate inverse Hessian),
+and its Hessian-vector product (``hessian_vector_product``, a directional
+derivative of the analytic force computed inside VMEC++).
+
+This module wires that residual to several solvers and is shared by the
+benchmark ``main`` below and by the tests:
+
+* preconditioned descent (VMEC's own update direction),
+* Jacobian-free Newton-Krylov, plain and preconditioned, and
+* a true Newton-Krylov driven by VMEC++'s own Hessian-vector product.
+
+All converge to the same equilibrium as the native solver. Force evaluations are
+counted inside VMEC++ (``force_eval_count``) so the comparison is fair across
+methods, including the evaluations hidden in Hessian-vector products.
+
+Interpretation: what separates the methods is how they couple DOFs of
+neighbouring flux surfaces. Plain JFNK does not -- it treats every degree of
+freedom as independent, so it cannot damp the radial stiffness that lets a
+surface cross its neighbour between iterations (BAD_JACOBIAN restarts in the
+native solver). It still converges, but slowly. VMEC's couples each surface to
+its jF +/- 1 neighbours; applying it -- as the inner Krylov preconditioner, or
+refreshed every step in the HVP Newton -- is what rescues conditioning and cuts
+the cost by an order of magnitude.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+from scipy.optimize import newton_krylov
+from scipy.sparse.linalg import LinearOperator, gmres
+
+import vmecpp
+from vmecpp.cpp import _vmecpp  # type: ignore[import]
+
+DEFAULT_INPUT = (
+    Path(__file__).resolve().parents[1]
+    / "examples"
+    / "data"
+    / "cth_like_fixed_bdy.json"
+)
+
+
+_BAD_JACOBIAN = 2  # RestartReason::BAD_JACOBIAN (flow_control.h)
+
+
+def _ensure_valid_initial_jacobian(model, max_reguess: int = 2) -> None:
+    """Reguess the magnetic axis until the initial Jacobian is non-singular.
+
+    Inputs that ship no axis (e.g. cma.json, with raxis/zaxis all zero) start from a
+    degenerate geometry, so the bare forward model returns at the BAD_JACOBIAN
+    checkpoint with zero force. The native solver reguesses the axis on the first
+    iterate (vmec.cc SolveEquilibriumLoop); mirror that here via reinitialize() so a
+    cold start from any valid INDATA has a defined gradient.
+    """
+    for _ in range(max_reguess):
+        model.evaluate(2, 2, False)
+        if model.restart_reason != _BAD_JACOBIAN:
+            return
+        model.reinitialize()
+    model.evaluate(2, 2, False)
+
+
+def make_model(input_path: Path = DEFAULT_INPUT, ns: int = 11):
+    indata = _vmecpp.VmecINDATA.from_file(str(input_path))
+    model = _vmecpp.VmecModel.create(indata, ns)
+    _ensure_valid_initial_jacobian(model)
+    return model
+
+
+def residual(model):
+    """Return F(x) = raw internal-basis force; F(x) = 0 at equilibrium."""
+
+    def F(x):
+        model.set_state(np.ascontiguousarray(x, dtype=float))
+        model.evaluate(2, 2, False)
+        return np.asarray(model.get_forces(), dtype=float)
+
+    return F
+
+
+@dataclass
+class Result:
+    name: str
+    force_evals: int
+    outer_iters: int
+    seconds: float
+    residual_norm: float
+    energy: float
+
+
+def reference_equilibrium(input_path: Path = DEFAULT_INPUT, ns: int = 11):
+    model = make_model(input_path, ns)
+    model.solve()
+    model.evaluate(2, 2, False)
+    return np.asarray(model.get_state(), float), model.mhd_energy
+
+
+def _finish(model, name, x, outer_iters, t0):
+    model.set_state(np.ascontiguousarray(x))
+    model.evaluate(2, 2, False)
+    return x, Result(
+        name,
+        model.force_eval_count,
+        outer_iters,
+        time.perf_counter() - t0,
+        float(np.linalg.norm(np.asarray(model.get_forces(), float))),
+        model.mhd_energy,
+    )
+
+
+def solve_vmecpp(input_path=DEFAULT_INPUT, ns=11):
+    """Native VMEC++ solve through the public API; reports the iteration count.
+
+    Unlike the other variants this drives the ordinary ``vmecpp.run`` path rather
+    than the low-level ``VmecModel`` primitives, so the reported residual/energy
+    are the public wout diagnostics (invariant force residuals and total MHD
+    energy) rather than the raw internal-basis force this module counts elsewhere.
+    """
+    vmec_input = vmecpp.VmecInput.from_file(input_path)
+    vmec_input.ns_array = np.asarray([ns], dtype=vmec_input.ns_array.dtype)
+    t0 = time.perf_counter()
+    output = vmecpp.run(vmec_input, verbose=False)
+    wout = output.wout
+    return output, Result(
+        name="VMEC++ (native, vmecpp.run)",
+        force_evals=wout.niter,
+        outer_iters=wout.niter,
+        seconds=time.perf_counter() - t0,
+        residual_norm=float(np.sqrt(wout.fsqt[-1])),
+        energy=wout.wb + wout.wp,
+    )
+
+
+def solve_preconditioned_descent(
+    input_path=DEFAULT_INPUT, ns=11, tol=1e-9, delt=0.9, momentum=0.5, max_iter=20000
+):
+    model = make_model(input_path, ns)
+    F = residual(model)
+    x = np.asarray(model.get_state(), float).copy()
+    v = np.zeros_like(x)
+    model.reset_force_eval_count()
+    it = 0
+    t0 = time.perf_counter()
+    for _ in range(max_iter):
+        if np.linalg.norm(F(x)) < tol:
+            break
+        it += 1
+        model.set_state(np.ascontiguousarray(x))
+        model.evaluate(2, 2, True)  # preconditioned search direction
+        fprec = np.asarray(model.get_forces(), float)
+        v = momentum * v + delt * fprec
+        x = x + delt * v
+    return _finish(model, "preconditioned descent", x, it, t0)
+
+
+def solve_newton_krylov(
+    input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=2500, preconditioned=False
+):
+    model = make_model(input_path, ns)
+    F = residual(model)
+    x0 = np.asarray(model.get_state(), float)
+    inner_m = None
+    model.reset_force_eval_count()
+    if preconditioned:
+        # Assemble VMEC's preconditioner at x0 and use it, frozen, as the inner
+        # Krylov preconditioner. M^-1 approximates the inverse Hessian; it is a
+        # radial tridiagonal (Thomas) solve per Fourier mode, so it couples each
+        # flux surface to its jF +/- 1 neighbours -- the coupling the plain
+        # branch lacks.
+        model.evaluate(2, 2, True)
+        n_dof = x0.size
+        inner_m = LinearOperator(  # type: ignore[call-overload]
+            (n_dof, n_dof),
+            matvec=lambda b: np.asarray(  # type: ignore[call-overload]
+                model.apply_preconditioner(np.ascontiguousarray(b)), float
+            ),
+        )
+    t0 = time.perf_counter()
+    x = newton_krylov(
+        F, x0, f_tol=tol, maxiter=max_iter, method="lgmres", inner_M=inner_m
+    )
+    name = (
+        "Newton-Krylov (preconditioned)" if preconditioned else "Newton-Krylov (JFNK)"
+    )
+    return _finish(model, name, x, 0, t0)
+
+
+def solve_newton_krylov_preconditioned(input_path=DEFAULT_INPUT, ns=11, tol=1e-9):
+    return solve_newton_krylov(input_path, ns, tol, preconditioned=True)
+
+
+def solve_newton_hvp(
+    input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_newton=80, inner_tol=1e-3
+):
+    """Globalized Newton-Krylov using VMEC++'s own Hessian-vector product.
+
+    Each Newton step solves H dx = -F with GMRES, where H v is hessian_vector_product
+    (the analytic force's directional derivative computed inside VMEC++) and the inner
+    solve is preconditioned by M^-1. A backtracking line search on ||F|| globalizes the
+    step, which is required on stiff 3D cases where the full Newton step overshoots.
+    """
+    model = make_model(input_path, ns)
+    F = residual(model)
+    x = np.asarray(model.get_state(), float).copy()
+    n_dof = x.size
+    model.reset_force_eval_count()
+    t0 = time.perf_counter()
+    it = 0
+    for _ in range(max_newton):
+        fk = F(x)
+        norm0 = np.linalg.norm(fk)
+        if norm0 < tol:
+            break
+        it += 1
+        model.set_state(np.ascontiguousarray(x))
+        model.evaluate(2, 2, True)  # assemble M at the current iterate
+        h_op = LinearOperator(  # type: ignore[call-overload]
+            (n_dof, n_dof),
+            matvec=lambda v: np.asarray(  # type: ignore[call-overload]
+                model.hessian_vector_product(np.ascontiguousarray(v)), float
+            ),
+        )
+        m_op = LinearOperator(  # type: ignore[call-overload]
+            (n_dof, n_dof),
+            matvec=lambda b: np.asarray(  # type: ignore[call-overload]
+                model.apply_preconditioner(np.ascontiguousarray(b)), float
+            ),
+        )
+        dx, _ = gmres(h_op, -fk, M=m_op, rtol=inner_tol, maxiter=100)
+        # Backtracking line search: accept the largest step that reduces ||F||.
+        alpha = 1.0
+        for _ in range(30):
+            if np.linalg.norm(F(x + alpha * dx)) < norm0:
+                break
+            alpha *= 0.5
+        else:
+            break  # no decrease found; stop
+        x = x + alpha * dx
+    return _finish(model, "Newton (VMEC++ HVP + M^-1)", x, it, t0)
+
+
+ALL_SOLVERS = (
+    solve_vmecpp,
+    solve_preconditioned_descent,
+    solve_newton_krylov_preconditioned,
+    solve_newton_krylov,
+    solve_newton_hvp,
+)
+
+
+def main():
+    _, w_star = reference_equilibrium()
+    print(f"reference equilibrium (native solve): W = {w_star:.8e}\n")
+    print(
+        f"{'optimizer':32s} {'F-evals':>8s} {'iters':>6s} {'time[s]':>8s} "
+        f"{'||F||':>10s} {'dW vs ref':>10s}"
+    )
+    for solver in ALL_SOLVERS:
+        r = solver()[1]
+        print(
+            f"{r.name:32s} {r.force_evals:8d} {r.outer_iters:6d} {r.seconds:8.2f} "
+            f"{r.residual_norm:10.1e} {abs(r.energy - w_star):10.1e}"
+        )
+
+
+if __name__ == "__main__":
+    main()

examples/fourier_resolution_increase.py

@@ -10,14 +10,13 @@
 resolution by :func:`vmecpp.interpolate_solution` (radial interpolation in
 ``sqrt(s)`` plus Fourier zero-padding).
 
-``vmecpp.run_continuation`` performs the whole schedule in one call:
+Setting ``VmecInput.mpol``/``.ntor`` to a sequence (instead of a plain int) makes
+``vmecpp.run`` perform the whole schedule in one call:
 
-    output = vmecpp.run_continuation(
-        vmec_input,
-        ns_array=[15, 31, 31],
-        mpol_array=[5, 9, 13],
-        ntor_array=[4, 4, 4],
-    )
+    vmec_input.ns_array = np.array([15, 31, 31])
+    vmec_input.mpol = [5, 9, 13]
+    vmec_input.ntor = [4, 4, 4]
+    output = vmecpp.run(vmec_input)
 
 This example runs the schedule step by step instead, so it can report the
 per-step iteration counts and compare the total against solving at the target
@@ -64,8 +63,8 @@ def step_input(ns: int, mpol: int, ntor: int) -> vmecpp.VmecInput:
 
 
 # --- Fourier continuation ---------------------------------------------------
-# vmecpp.run_continuation(vmec_input, ns_array=ns_array, mpol_array=mpol_array,
-# ntor_array=ntor_array) runs exactly this schedule in a single call; it is
+# Setting vmec_input.ns_array/.mpol/.ntor to these schedules and calling
+# vmecpp.run(vmec_input) runs exactly this schedule in a single call; it is
 # unrolled here to report the per-step iteration counts.
 print("Fourier continuation:")
 schedule = list(zip(ns_array, mpol_array, ntor_array, strict=True))