Hand every force iteration to a callback and watch a solve converge

README.md

@@ -101,6 +101,22 @@ vmec_output.wout.save("wout_w7x.nc")
 
 All other output files are accessible via members of the `vmec_output` object called `threed1_volumetrics`, `jxbout` and `mercier`.
 
+### Watching a solve
+
+`vmecpp.watch` runs the same solve as `vmecpp.run` and draws the flux surfaces and the
+force residuals of every iteration while it runs, in a window or into an animation file:
+
+```python
+import vmecpp
+
+vmec_input = vmecpp.VmecInput.from_file("examples/data/solovev.json")
+output = vmecpp.watch(vmec_input, save="solve.gif")
+```
+
+The per-iteration data behind it is available to any script through the
+`iteration_callback` argument of `vmecpp.run`, which receives an `IterationSnapshot` with
+the force residuals, the flow-control state and the geometry of every iteration.
+
 ### With SIMSOPT
 
 [SIMSOPT](https://simsopt.readthedocs.io) is a popular stellarator optimization framework.

examples/watch_solve.py

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH <info@proximafusion.com>
+#
+# SPDX-License-Identifier: MIT
+"""Watch VMEC++ converge: the flux surfaces at two toroidal angles above the force
+residuals of every iteration, drawn while the solve runs.
+
+Run it with a display to get a live window, or pass ``--save solve.gif`` to record
+the solve headlessly. The same per-iteration data is available to any script
+through the ``iteration_callback`` argument of ``vmecpp.run``.
+"""
+
+import argparse
+from pathlib import Path
+
+import matplotlib as mpl
+from matplotlib.backends.registry import BackendFilter, backend_registry
+
+import vmecpp
+
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument(
+    "input",
+    nargs="?",
+    default=Path(__file__).parent / "data" / "solovev.json",
+    type=Path,
+    help="a VMEC++ JSON or INDATA file (default: examples/data/solovev.json)",
+)
+parser.add_argument(
+    "--save",
+    type=Path,
+    help="record the solve to a .gif or video file (the default without a display)",
+)
+parser.add_argument("--every", type=int, default=5, help="draw every N iterations")
+parser.add_argument(
+    "--planes", type=int, default=2, help="toroidal cross-sections to show"
+)
+args = parser.parse_args()
+
+interactive = {
+    name.lower() for name in backend_registry.list_builtin(BackendFilter.INTERACTIVE)
+}
+save = args.save
+if save is None and mpl.get_backend().lower() not in interactive:
+    save = Path("watch_solve.gif")
+    print(f"no display, recording the solve to {save}")
+
+vmec_input = vmecpp.VmecInput.from_file(args.input)
+output = vmecpp.watch(vmec_input, planes=args.planes, every=args.every, save=save)
+print(
+    f"{output.wout.reason}: fsqr = {output.wout.fsqr:.2e} after {output.wout.niter} iterations"
+)

src/vmecpp/__init__.py

@@ -40,6 +40,7 @@
     own_model_fields,
 )
 from vmecpp._rescale import rescale
+from vmecpp._watch import watch
 from vmecpp.cpp import _vmecpp  # type: ignore # bindings to the C++ core
 
 logger = logging.getLogger(__name__)
@@ -862,6 +863,7 @@ def reason(self) -> str:
         return {
             0: "normal termination: converged, or returned without convergence because return_outputs_even_if_not_converged was set",
             1: "initially bad Jacobian",
+            2: "stopped by the iteration callback before convergence",
             3: "NCURR_NE_1_BLOAT_NE_1",
             4: "Jacobian reset 75 times, the geometry isn't well defined",
             5: "unrecoverable error: a physical inconsistency in the MHD model, such as a degenerate flux-surface geometry or a free-boundary current mismatch, with no retry strategy",
@@ -2411,13 +2413,29 @@ def _print_progress_tip_once() -> None:
         )
 
 
+IterationSnapshot: typing.TypeAlias = _vmecpp.IterationSnapshot
+"""The state of one force iteration, handed to the ``iteration_callback`` of
+:func:`run`.
+
+Its attributes are ``iteration`` (the counter of the current multigrid stage, as
+printed), ``multigrid_step`` (the index into ``ns_array``, -1 for the inserted
+ns = 3 stage), ``ns``, the invariant force residuals ``fsqr``, ``fsqz``, ``fsql``
+and the ``ftol`` they are tested against, the time step ``delt``,
+``restart_reason`` (1 no restart, 2 bad Jacobian, 3 bad progress, 4 huge initial
+forces), ``jacobian_resets``, ``vacuum_pressure_active``, ``mhd_energy`` and
+``geometry``, the R, Z and lambda coefficients of the state as a ``Geometry`` that
+``vmecpp.geometry.from_cpp`` reads.
+"""
+
+
 def run(
     input: VmecInput,
     magnetic_field: MagneticFieldResponseTable | None = None,
     *,
     max_threads: int | None = None,
     verbose: bool | int | OutputMode = OutputMode.PROGRESS,
     restart_from: VmecOutput | None = None,
+    iteration_callback: typing.Callable[[IterationSnapshot], bool | None] | None = None,
 ) -> VmecOutput:
     """Run VMEC++ using the provided input. This is the main entrypoint for both fixed-
     and free-boundary calculations.
@@ -2439,6 +2457,11 @@ def run(
             convergence when running VMEC++ on a configuration that is very similar to the `restart_from` equilibrium.
             If `input.mpol`/`input.ntor` is a sequence (see below), this is used to hot-restart
             only the first continuation step; later steps always hot-restart from the previous one.
+        iteration_callback: called once per force iteration with an :class:`IterationSnapshot`
+            of the state just reached, after every thread has finished the step. Returning
+            ``False`` stops the run, which then returns the outputs of that state with
+            ``wout.ier_flag`` reporting no convergence; returning ``None`` or ``True`` continues.
+            An exception raised inside the callback stops the run and propagates.
 
     If `input.mpol` and/or `input.ntor` is a sequence rather than a plain int, `run` performs
     continuation in Fourier resolution: each entry pairs with the corresponding `input.ns_array`
@@ -2463,6 +2486,7 @@ def run(
             max_threads=max_threads,
             verbose=verbose,
             restart_from=restart_from,
+            iteration_callback=iteration_callback,
         )
 
     cpp_indata = input._to_cpp_vmecindata()
@@ -2499,6 +2523,7 @@ def run(
             initial_state=initial_state,
             max_threads=max_threads,
             verbose=_verbose.value,
+            iteration_callback=iteration_callback,
         )
     else:
         # magnetic_response_table takes precedence anyway, but let's be explicit, to ensure
@@ -2510,6 +2535,7 @@ def run(
             initial_state=initial_state,
             max_threads=max_threads,
             verbose=_verbose.value,
+            iteration_callback=iteration_callback,
         )
 
     cpp_wout = cpp_output_quantities.wout
@@ -2729,4 +2755,6 @@ def set_profile(
     "solve_multigrid",
     "IterationResult",
     "IterationState",
+    "IterationSnapshot",
+    "watch",
 ]

src/vmecpp/_continuation.py

@@ -31,7 +31,7 @@
 import numpy as np
 
 if typing.TYPE_CHECKING:
-    from vmecpp import OutputMode, VmecInput, VmecOutput
+    from vmecpp import IterationSnapshot, OutputMode, VmecInput, VmecOutput
     from vmecpp._free_boundary import MagneticFieldResponseTable
 
 # State-vector geometry arrays, shape [mn_mode, n_surfaces]. These are the only
@@ -301,6 +301,7 @@ def _run_fourier_continuation(
     max_threads: int | None,
     verbose: bool | int | OutputMode,
     restart_from: VmecOutput | None,
+    iteration_callback: typing.Callable[[IterationSnapshot], bool | None] | None = None,
 ) -> VmecOutput:
     """Solves an equilibrium by continuation in Fourier resolution.
 
@@ -363,6 +364,7 @@ def _resolve(value: int | np.ndarray, name: str) -> list[int]:
             max_threads=max_threads,
             verbose=verbose,
             restart_from=guess,
+            iteration_callback=iteration_callback,
         )
 
     assert output is not None  # n_steps >= 1, so the loop always assigns output