Add a geometry tolerance that holds a run until the flux surfaces stop moving

src/vmecpp/__init__.py

@@ -446,6 +446,16 @@ class VmecInput(BaseModelWithNumpy):
     tcon0: float = 1.0
     """Constraint force scaling factor for ns --> 0."""
 
+    geometry_tolerance: float = 0.0
+    """Distance in metres below which the geometry counts as settled, measured over the
+    last ``nstep`` iterations as the Euclidean norm of the change in the R and Z
+    spectral coefficients.
+
+    When positive, a multigrid step converges only once the force residuals meet
+    ``ftol_array`` and the geometry has moved less than this. Zero leaves convergence on
+    the residuals alone.
+    """
+
     lforbal: bool = False
     """Hack: directly compute innermost flux surface geometry from radial force balance"""
 

src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h

@@ -98,6 +98,12 @@ class FlowControl {
   double fsqr1, fsqz1, fsql1;
   double fsq;
 
+  // How far the boundary-normalized geometry moved over the last `nstep`
+  // iterations, in metres: the Euclidean norm of the change in the R and Z
+  // spectral coefficients. Populated by `Vmec::AccumulateGeometryChange` at
+  // each printout, and negative until two printouts have happened.
+  double geometry_change = -1.0;
+
   std::vector<double> mhd_energy;
 
   // Time-trace of the force at the vacuum boundary (only for free-boundary)

src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc

@@ -227,6 +227,7 @@ VmecINDATA::VmecINDATA() {
   aphi[0] = 1.0;
   delt = 1.0;
   tcon0 = 1.0;
+  geometry_tolerance = 0.0;
   lforbal = false;
   iteration_style = IterationStyle::VMEC_8_52;
   return_outputs_even_if_not_converged = false;
@@ -353,6 +354,7 @@ absl::Status VmecINDATA::WriteTo(H5::H5File& file) const {
   WriteH5Dataset(nstep, "/indata/nstep", file);
   WriteH5Dataset(delt, "/indata/delt", file);
   WriteH5Dataset(tcon0, "/indata/tcon0", file);
+  WriteH5Dataset(geometry_tolerance, "/indata/geometry_tolerance", file);
   WriteH5Dataset(lforbal, "/indata/lforbal", file);
   WriteH5Dataset(return_outputs_even_if_not_converged,
                  "/indata/return_outputs_even_if_not_converged", file);
@@ -458,6 +460,10 @@ absl::Status VmecINDATA::LoadInto(VmecINDATA& m_indata, H5::H5File& from_file) {
   ReadH5Dataset(m_indata.nstep, "/indata/nstep", from_file);
   ReadH5Dataset(m_indata.delt, "/indata/delt", from_file);
   ReadH5Dataset(m_indata.tcon0, "/indata/tcon0", from_file);
+  if (from_file.nameExists("/indata/geometry_tolerance")) {
+    ReadH5Dataset(m_indata.geometry_tolerance, "/indata/geometry_tolerance",
+                  from_file);
+  }
   ReadH5Dataset(m_indata.lforbal, "/indata/lforbal", from_file);
 
   // Legacy way of checking for dataset existence
@@ -935,6 +941,14 @@ absl::StatusOr<VmecINDATA> VmecINDATA::FromJson(
     vmec_indata.tcon0 = maybe_tcon0->value();
   }
 
+  auto maybe_geometry_tolerance = JsonReadDouble(j, "geometry_tolerance");
+  if (!maybe_geometry_tolerance.ok()) {
+    return maybe_geometry_tolerance.status();
+  }
+  if (maybe_geometry_tolerance->has_value()) {
+    vmec_indata.geometry_tolerance = maybe_geometry_tolerance->value();
+  }
+
   auto maybe_lforbal = JsonReadBool(j, "lforbal");
   if (!maybe_lforbal.ok()) {
     return maybe_lforbal.status();
@@ -1266,6 +1280,7 @@ absl::StatusOr<std::string> VmecINDATA::ToJson() const {
   output["aphi"] = aphi;
   output["delt"] = delt;
   output["tcon0"] = tcon0;
+  output["geometry_tolerance"] = geometry_tolerance;
   output["lforbal"] = lforbal;
   output["iteration_style"] = ToString(iteration_style);
   output["return_outputs_even_if_not_converged"] =
@@ -1575,6 +1590,13 @@ absl::Status IsConsistent(const VmecINDATA& vmec_indata,
         vmec_indata.delt));
   }
 
+  if (vmec_indata.geometry_tolerance < 0.0) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "input variable 'geometry_tolerance' is a distance and cannot be "
+        "negative, but is %g\n",
+        vmec_indata.geometry_tolerance));
+  }
+
   // tcon0
   if (vmec_indata.tcon0 < 0.0 || vmec_indata.tcon0 > 1.0) {
     return absl::InvalidArgumentError(absl::StrFormat(

src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.h

@@ -215,6 +215,13 @@ class VmecINDATA {
   // constraint force scaling factor for ns --> 0
   double tcon0;
 
+  // Distance in metres below which the geometry counts as settled, measured
+  // over the last nstep iterations as the Euclidean norm of the change in the
+  // R and Z spectral coefficients. When positive, a multigrid step converges
+  // only once the force residuals meet ftol AND the geometry has moved less
+  // than this. Zero leaves convergence on the residuals alone.
+  double geometry_tolerance;
+
   // hack: directly compute innermost flux surface geometry from radial force
   // balance
   bool lforbal;

src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.cc

@@ -4,6 +4,7 @@
 // SPDX-License-Identifier: MIT
 #include "vmecpp/vmec/handover_storage/handover_storage.h"
 
+#include <cmath>
 #include <iostream>
 
 namespace vmecpp {
@@ -39,6 +40,7 @@ HandoverStorage::HandoverStorage(const Sizes* s) : s_(*s) {
   // as a division-by-zero would occur.
   spectral_width_numerator_ = 0.0;
   spectral_width_denominator_ = 0.0;
+  geometry_change_ = 0.0;
 
   rAxis.setZero(s_.nZeta);
   zAxis.setZero(s_.nZeta);
@@ -190,6 +192,10 @@ double HandoverStorage::VolumeAveragedSpectralWidth() const {
   return spectral_width_numerator_ / spectral_width_denominator_;
 }  // VolumeAveragedSpectralWidth
 
+double HandoverStorage::GeometryChange() const {
+  return std::sqrt(geometry_change_);
+}  // GeometryChange
+
 void HandoverStorage::SetRadialExtent(const RadialExtent& radial_extent) {
   radial_extent_ = radial_extent;
 }  // SetRadialExtent

src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.h

@@ -51,6 +51,11 @@ class HandoverStorage {
   double* SpectralWidthDenominator() { return &spectral_width_denominator_; }
   double VolumeAveragedSpectralWidth() const;
 
+  // Destination of the cross-thread fold of the squared geometry change, and
+  // the distance that fold amounts to.
+  double* GeometryChangeAccumulator() { return &geometry_change_; }
+  double GeometryChange() const;
+
   void SetRadialExtent(const RadialExtent& radial_extent);
   void SetGeometricOffset(const GeometricOffset& geometric_offset);
 
@@ -220,6 +225,7 @@ class HandoverStorage {
 
   double spectral_width_numerator_;
   double spectral_width_denominator_;
+  double geometry_change_;
 
   RadialExtent radial_extent_;
   GeometricOffset geometric_offset_;

src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc

@@ -832,6 +832,7 @@ PYBIND11_MODULE(_vmecpp, m) {
   DefEigenProperty(pyindata, "aphi", &VmecINDATA::aphi);
   pyindata.def_readwrite("delt", &VmecINDATA::delt)
       .def_readwrite("tcon0", &VmecINDATA::tcon0)
+      .def_readwrite("geometry_tolerance", &VmecINDATA::geometry_tolerance)
       .def_readwrite("lforbal", &VmecINDATA::lforbal)
       .def_readwrite("iteration_style", &VmecINDATA::iteration_style)
       .def_readwrite("return_outputs_even_if_not_converged",

src/vmecpp/cpp/vmecpp/vmec/vmec/AGENTS.md

@@ -43,6 +43,12 @@ Reached when all three force residuals fall below the current stage tolerance:
 the iteration count exceeds `niterv`. Residuals live in `FlowControl` (`flow_control.h`);
 `fsq*1` are the preconditioned variants used for the damping average.
 
+The residual is small in directions the state is still moving along, so `indata.geometry_tolerance`
+adds a second condition when it is positive: `fc_.geometry_change`, the distance in metres the R
+and Z coefficients moved over the last `nstep` iterations, must also be below it. It is folded
+across the team in `Vmec::AccumulateGeometryChange()` at each printout and reset at every
+multigrid stage, since the radial grid changes between stages.
+
 ## Restart logic
 
 `Vmec::RestartIteration()` (enum `RestartReason` in `flow_control.h`). When the iteration

src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc

@@ -349,6 +349,13 @@ absl::StatusOr<bool> Vmec::run(const VmecCheckpoint& checkpoint,
         fc_.restart_reasons.reserve(cap);
       }
 
+      // The radial grid changes between stages, so the geometry of the previous
+      // stage is not something this one can be compared against.
+      for (auto& geometry : geometry_at_last_printout_) {
+        geometry.reset();
+      }
+      fc_.geometry_change = -1.0;
+
       // notify logger of the next multigrid stage
       logger_.BeginStage(igrid, max_grids + jacob_off_, fc_.nsval, s_.mnmax,
                          fc_.ftolv, fc_.niterv, fc_.lfreeb);
@@ -624,6 +631,7 @@ absl::StatusOr<bool> Vmec::InitializeRadial(
     p_.resize(num_threads_);
     m_.resize(num_threads_);
     decomposed_x_.resize(num_threads_);
+    geometry_at_last_printout_.resize(num_threads_);
     physical_x_backup_.resize(num_threads_);
     physical_x_.resize(num_threads_);
     decomposed_f_.resize(num_threads_);
@@ -1363,8 +1371,12 @@ absl::StatusOr<bool> Vmec::Evolve(VmecCheckpoint checkpoint,
       // first iteration and Jacobian was not computed correctly
       status_ = VmecStatus::BAD_JACOBIAN;
     } else if (fc_.fsqr <= fc_.ftolv && fc_.fsqz <= fc_.ftolv &&
-               fc_.fsql <= fc_.ftolv) {
-      // converged to desired tolerance
+               fc_.fsql <= fc_.ftolv &&
+               (indata_.geometry_tolerance <= 0.0 ||
+                (fc_.geometry_change >= 0.0 &&
+                 fc_.geometry_change <= indata_.geometry_tolerance))) {
+      // converged to desired tolerance, and where a geometry tolerance is set,
+      // the flux surfaces have stopped moving as well
 
       m_liter_flag = false;
       status_ = VmecStatus::SUCCESSFUL_TERMINATION;
@@ -1452,6 +1464,51 @@ absl::StatusOr<bool> Vmec::Evolve(VmecCheckpoint checkpoint,
   return false;
 }
 
+void Vmec::AccumulateGeometryChange(int thread_id) {
+  const FourierGeometry& x = *decomposed_x_[thread_id];
+  const RadialPartitioning& r = *r_[thread_id];
+  const int mnsize = s_.mnsize;
+  const int offset = (r.nsMinF - x.nsMin()) * mnsize;
+  const int count = (r.nsMaxFIncludingLcfs - r.nsMinF) * mnsize;
+
+  // Surfaces [nsMinF, nsMaxFIncludingLcfs) partition the plasma across the
+  // team, so no surface is counted twice and the measure does not depend on
+  // how many threads run.
+  const std::array<std::span<double>, 8> now = {
+      x.rmncc, x.rmnss, x.rmnsc, x.rmncs, x.zmnsc, x.zmncs, x.zmncc, x.zmnss};
+  double contribution = 0.0;
+  const FourierGeometry* previous = geometry_at_last_printout_[thread_id].get();
+  if (previous != nullptr) {
+    const std::array<std::span<double>, 8> before = {
+        previous->rmncc, previous->rmnss, previous->rmnsc, previous->rmncs,
+        previous->zmnsc, previous->zmncs, previous->zmncc, previous->zmnss};
+    for (size_t block = 0; block < now.size(); ++block) {
+      if (now[block].size() != before[block].size() ||
+          static_cast<int>(now[block].size()) < offset + count) {
+        continue;
+      }
+      for (int i = offset; i < offset + count; ++i) {
+        const double difference = now[block][i] - before[block][i];
+        contribution += difference * difference;
+      }
+    }
+  }
+
+  SumOverThreads(&contribution, 1, thread_id, r.get_num_threads(),
+                 h_.thread_reduce_slots.data(), h_.GeometryChangeAccumulator());
+
+#ifdef _OPENMP
+#pragma omp single nowait
+#endif  // _OPENMP
+  {
+    // Every thread creates its copy in the same pass, so this test is the same
+    // on all of them and it does not matter which one runs the block.
+    fc_.geometry_change = previous == nullptr ? -1.0 : h_.GeometryChange();
+  }
+
+  geometry_at_last_printout_[thread_id] = std::make_unique<FourierGeometry>(x);
+}
+
 void Vmec::Printout(double delt0r, int thread_id, int iter2) {
 #ifdef _OPENMP
 #pragma omp single
@@ -1460,6 +1517,7 @@ void Vmec::Printout(double delt0r, int thread_id, int iter2) {
     h_.ResetSpectralWidthAccumulators();
   }
   p_[thread_id]->AccumulateVolumeAveragedSpectralWidth();
+  AccumulateGeometryChange(thread_id);
 #ifdef _OPENMP
 #pragma omp barrier
 #endif  // _OPENMP

src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h

@@ -146,6 +146,11 @@ class Vmec {
                               double time_step, int thread_id,
                               bool& m_liter_flag);
   void Printout(double delt0r, int thread_id, int iter2);
+
+  // Fold the squared change of this thread's R and Z coefficients since the
+  // previous printout into fc_.geometry_change, then keep the current geometry
+  // for the next one. Every thread of the team must call this.
+  void AccumulateGeometryChange(int thread_id);
   absl::StatusOr<bool> UpdateForwardModel(VmecCheckpoint checkpoint,
                                           int maximum_iterations,
                                           int thread_id);
@@ -220,6 +225,11 @@ class Vmec {
   std::vector<std::unique_ptr<FourierForces>> physical_f_;
   std::vector<std::unique_ptr<FourierVelocity>> decomposed_v_;
 
+  // Geometry as it stood at the previous printout, per thread, against which
+  // fc_.geometry_change is measured. Empty until the first printout of a
+  // multigrid step, since the radial resolution changes between steps.
+  std::vector<std::unique_ptr<FourierGeometry>> geometry_at_last_printout_;
+
   std::vector<std::unique_ptr<FourierGeometry>> old_xc_scaled_;
   std::vector<std::unique_ptr<RadialPartitioning>> old_r_;