From b237876c08b4ba7d189cffa880da7428b7474c7d Mon Sep 17 00:00:00 2001 From: CharlesCNorton <135471798+CharlesCNorton@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:45:23 -0400 Subject: [PATCH 1/2] Add a geometry tolerance that holds a run until the flux surfaces stop moving --- src/vmecpp/__init__.py | 10 ++ .../vmecpp/common/flow_control/flow_control.h | 6 + .../vmecpp/common/vmec_indata/vmec_indata.cc | 22 ++++ .../vmecpp/common/vmec_indata/vmec_indata.h | 7 ++ .../vmec/handover_storage/handover_storage.cc | 6 + .../vmec/handover_storage/handover_storage.h | 6 + .../cpp/vmecpp/vmec/pybind11/pybind_vmec.cc | 1 + src/vmecpp/cpp/vmecpp/vmec/vmec/AGENTS.md | 6 + src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc | 67 +++++++++- src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h | 10 ++ tests/test_geometry_tolerance.py | 115 ++++++++++++++++++ 11 files changed, 254 insertions(+), 2 deletions(-) create mode 100644 tests/test_geometry_tolerance.py diff --git a/src/vmecpp/__init__.py b/src/vmecpp/__init__.py index 7f6d65795..763f175c9 100644 --- a/src/vmecpp/__init__.py +++ b/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""" diff --git a/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h b/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h index ab7e2cb79..ddf8a9a9f 100644 --- a/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h +++ b/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 mhd_energy; // Time-trace of the force at the vacuum boundary (only for free-boundary) diff --git a/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc b/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc index bae883c5c..57bc69db6 100644 --- a/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc +++ b/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::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 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( diff --git a/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.h b/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.h index 4316d49f5..335368572 100644 --- a/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.h +++ b/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; diff --git a/src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.cc b/src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.cc index 7ab009676..fa328aa29 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.cc +++ b/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 #include 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 diff --git a/src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.h b/src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.h index fbbdd47bf..5e00ec018 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/handover_storage/handover_storage.h +++ b/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_; diff --git a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc index 675315791..7f6513cca 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc +++ b/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", diff --git a/src/vmecpp/cpp/vmecpp/vmec/vmec/AGENTS.md b/src/vmecpp/cpp/vmecpp/vmec/vmec/AGENTS.md index 861acee4b..cd718981c 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/AGENTS.md +++ b/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 diff --git a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc index a690ebaf6..2d6ccfb63 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc @@ -349,6 +349,13 @@ absl::StatusOr 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 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 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,56 @@ absl::StatusOr 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, 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, 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(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(); + } + + if (previous == nullptr) { + geometry_at_last_printout_[thread_id] = + std::make_unique(x); + } else { + *geometry_at_last_printout_[thread_id] = x; + } +} + void Vmec::Printout(double delt0r, int thread_id, int iter2) { #ifdef _OPENMP #pragma omp single @@ -1460,6 +1522,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 diff --git a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h index ed3485d1e..90dd727c1 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h +++ b/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 UpdateForwardModel(VmecCheckpoint checkpoint, int maximum_iterations, int thread_id); @@ -220,6 +225,11 @@ class Vmec { std::vector> physical_f_; std::vector> 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> geometry_at_last_printout_; + std::vector> old_xc_scaled_; std::vector> old_r_; diff --git a/tests/test_geometry_tolerance.py b/tests/test_geometry_tolerance.py new file mode 100644 index 000000000..833bc2735 --- /dev/null +++ b/tests/test_geometry_tolerance.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# SPDX-License-Identifier: MIT +"""Convergence on the state rather than on the force alone. + +The force residual is small in directions the equilibrium is still moving along, +so a run that meets ``ftol_array`` can leave the magnetic axis and the rotational +transform short of where they settle. ``geometry_tolerance`` adds the condition +that the flux surfaces have stopped moving, measured over the last ``nstep`` +iterations as the Euclidean norm of the change in the R and Z coefficients. +""" + +from pathlib import Path + +import numpy as np +import pytest + +import vmecpp + +REPO_ROOT = Path(__file__).parent.parent +TEST_DATA = REPO_ROOT / "src" / "vmecpp" / "cpp" / "vmecpp" / "test_data" +CASE = TEST_DATA / "cth_like_fixed_bdy.json" + + +def _single_grid(ns: int = 51) -> vmecpp.VmecInput: + vmec_input = vmecpp.VmecInput.from_file(CASE) + vmec_input.ns_array = np.array([ns]) + vmec_input.niter_array = np.array([100000]) + return vmec_input + + +def _axis_and_iota(wout) -> tuple[float, float]: + return float(wout.rmnc[0, 0]), float(wout.iotaf[wout.iotaf.size // 2]) + + +def test_unset_tolerance_leaves_the_run_untouched(): + """The default is off, so the run is the one the residual test alone gives.""" + vmec_input = _single_grid() + vmec_input.ftol_array = np.array([1.0e-8]) + assert vmec_input.geometry_tolerance == 0.0 + loose = vmecpp.run(vmec_input, verbose=False, max_threads=1) + + explicit = vmec_input.model_copy(update={"geometry_tolerance": 0.0}) + again = vmecpp.run(explicit, verbose=False, max_threads=1) + + assert again.wout.niter == loose.wout.niter + np.testing.assert_array_equal(again.wout.rmnc, loose.wout.rmnc) + + +def test_a_settled_geometry_is_required_before_convergence(): + """At a tolerance most inputs ship with, the axis and iota are still moving; + requiring the surfaces to settle recovers them.""" + reference = vmecpp.run( + _single_grid().model_copy(update={"ftol_array": np.array([1.0e-16])}), + verbose=False, + max_threads=1, + ) + axis_reference, iota_reference = _axis_and_iota(reference.wout) + + loose_input = _single_grid().model_copy(update={"ftol_array": np.array([1.0e-8])}) + loose = vmecpp.run(loose_input, verbose=False, max_threads=1) + axis_loose, iota_loose = _axis_and_iota(loose.wout) + + settled = vmecpp.run( + loose_input.model_copy(update={"geometry_tolerance": 1.0e-5}), + verbose=False, + max_threads=1, + ) + axis_settled, iota_settled = _axis_and_iota(settled.wout) + + # the residual alone stops with the axis a fraction of a millimetre out + assert abs(axis_loose - axis_reference) > 1.0e-4 + assert abs(iota_loose - iota_reference) > 1.0e-4 + + # and the same run, held until the surfaces stop moving, lands on it + assert settled.wout.niter > loose.wout.niter + assert abs(axis_settled - axis_reference) < 1.0e-6 + assert abs(iota_settled - iota_reference) < 1.0e-5 + + +def test_a_tighter_tolerance_gets_closer(): + reference = vmecpp.run( + _single_grid().model_copy(update={"ftol_array": np.array([1.0e-16])}), + verbose=False, + max_threads=1, + ) + _, iota_reference = _axis_and_iota(reference.wout) + + errors = [] + for tolerance in (1.0e-5, 1.0e-6): + output = vmecpp.run( + _single_grid().model_copy( + update={ + "ftol_array": np.array([1.0e-8]), + "geometry_tolerance": tolerance, + } + ), + verbose=False, + max_threads=1, + ) + errors.append(abs(_axis_and_iota(output.wout)[1] - iota_reference)) + assert errors[1] < errors[0] + + +def test_a_negative_tolerance_is_rejected(): + vmec_input = _single_grid().model_copy(update={"geometry_tolerance": -1.0}) + with pytest.raises(AttributeError, match="geometry_tolerance"): + vmecpp.run(vmec_input, verbose=False, max_threads=1) + + +def test_the_tolerance_survives_a_round_trip(tmp_path): + vmec_input = _single_grid().model_copy(update={"geometry_tolerance": 2.5e-6}) + path = tmp_path / "input.json" + path.write_text(vmec_input.to_json()) + assert vmecpp.VmecInput.from_file(path).geometry_tolerance == 2.5e-6 From 39358da450e2c3cba873637f5507cd2953eab75b Mon Sep 17 00:00:00 2001 From: CharlesCNorton <135471798+CharlesCNorton@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:56:11 -0400 Subject: [PATCH 2/2] Take a fresh copy of the geometry at each printout instead of branching on the first one --- src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc index 2d6ccfb63..67cf236a8 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc @@ -1506,12 +1506,7 @@ void Vmec::AccumulateGeometryChange(int thread_id) { fc_.geometry_change = previous == nullptr ? -1.0 : h_.GeometryChange(); } - if (previous == nullptr) { - geometry_at_last_printout_[thread_id] = - std::make_unique(x); - } else { - *geometry_at_last_printout_[thread_id] = x; - } + geometry_at_last_printout_[thread_id] = std::make_unique(x); } void Vmec::Printout(double delt0r, int thread_id, int iter2) {