From 2449ddc9c59d0a4cbc676fbe30b0dc1dfdc72089 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Thu, 9 Jul 2026 14:55:50 +0200 Subject: [PATCH 01/12] pybind: Hessian-vector product inside VMEC++ + internal Newton-Krylov (#580) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * build: bump CMake abseil pin to 20260107.1 for Clang >= 21 The CMake FetchContent abseil pin (2024-08) fails to compile under Clang >= 21: absl::Nonnull SFINAE in absl/strings/ascii.cc and the numbers.cc nullability annotations are rejected by the newer frontend. Bump to the 20260107.1 LTS, which compiles cleanly under Clang 21.1.8 and GCC. Clang is the compiler required for the Enzyme autodiff build. The Bazel build keeps its own (BCR) abseil pin and is unaffected. * enzyme: opt-in Clang/Enzyme build option and AD smoke test Add VMECPP_ENABLE_ENZYME (OFF by default), which requires a Clang compiler and a ClangEnzyme plugin path and builds a self-contained autodiff smoke test. The test differentiates a scalar objective written over Eigen::Map'd caller buffers and checks reverse- and forward-mode Enzyme gradients against the closed form and central finite differences. enzyme.h documents the intrinsic ABI and the allocation constraint that shapes the differentiable kernels: Enzyme cannot track Eigen's aligned allocator, so differentiable paths use Eigen::Map over caller-owned buffers and avoid heap expression temporaries. With the option off the build is unchanged. * pybind: expose the unpreconditioned internal-basis gradient Add a precondition flag to VmecModel.evaluate (default true, unchanged behaviour). With precondition=false the forward model returns at the INVARIANT_RESIDUALS checkpoint, so get_forces() yields the raw, unpreconditioned force: the gradient of VMEC's augmented functional (MHD energy plus the spectral-condensation and lambda constraints) with respect to the decomposed internal-basis state. This is the consistent state/gradient pair an external optimizer needs to minimise in VMEC's own basis. The native solver's preconditioned search direction (precondition=true) is a different vector; the raw gradient is the equilibrium residual and vanishes at convergence. Tests: raw force is finite and differs in direction from the preconditioned force, and drops by >1e6 from the initial guess to the converged equilibrium. * examples: drive VMEC++ from external optimizers in the internal basis Treat the equilibrium as the root problem F(x) = 0, where F is the raw internal-basis force (gradient of VMEC's augmented functional) exposed by evaluate(precondition=False). Wire it to two solvers that reuse VMEC++'s forward model: native-style preconditioned descent and Jacobian-free Newton-Krylov (matrix-free Hessian information). Both reach the native solver's equilibrium. This is the external-differentiability path: VMEC++ as a differentiable equilibrium component an outside optimizer can drive. Quasi-Newton root-finders without a preconditioner diverge on this stiff system, which motivates exposing VMEC's preconditioner as an operator next. Tests assert both solvers reach force balance and recover the native energy and state. * pybind: expose VMEC preconditioner as an operator; preconditioned JFNK Add VmecModel.apply_preconditioner(v): applies VMEC's preconditioner M^-1 (m=1, radial, lambda steps) to a vector in the decomposed basis. M^-1 is VMEC's hand-built approximate inverse Hessian; this exposes it as a reusable linear operator for preconditioned Krylov / quasi-Newton and for the Hessian solve in adjoint sensitivities. It requires a prior evaluate(precondition=true), which assembles the radial preconditioner. Validated exactly: apply_preconditioner(raw force) equals the native preconditioned search direction; the operator is linear and, once assembled, state-invariant. Use it as the inner Krylov preconditioner in Newton-Krylov: on solovev (ns=11) this cuts force evaluations from 2242 to 505 (4.4x) versus unpreconditioned JFNK, converging to the same equilibrium. * pybind: Hessian-vector product inside VMEC++; internal Newton-Krylov Add VmecModel.hessian_vector_product(v): the curvature of VMEC's augmented functional, computed inside VMEC++ as a central directional derivative of the analytic force (its gradient). The force is exact; only the directional step is finite-differenced. Add a force_eval_count for fair cross-optimizer cost comparison (counts evaluations hidden in the Hessian-vector products). Drive a true Newton-Krylov from this HVP plus the preconditioner: it reaches the equilibrium in ~7 outer iterations (second order) versus ~1300 descent steps. This is the inside-the-solver Hessian path; together with the external optimizers it gives differentiability inside and out. Benchmark (solovev, ns=11, force evals counted in VMEC++): preconditioned descent 2606 evals 1302 iters Newton-Krylov (JFNK) 2243 evals Newton-Krylov (preconditioned) 507 evals Newton (VMEC++ HVP + M^-1) 9194 evals 7 iters The HVP-Newton's higher force-eval count (two evals per finite-difference HVP) is what the exact Enzyme Hessian will remove. * ideal_mhd_model: make computeMHDForces allocation-free The force kernel allocated 17 dynamic Eigen vectors per radial surface (the _o half-grid quantities and the avg/wavg surface averages). Move them to preallocated per-thread ThreadLocalStorage scratch and assign in place, so the radial loop allocates nothing. Two benefits: it removes per-surface heap churn from the hot force loop, and it makes the kernel differentiable by Enzyme, which cannot trace dynamic Eigen temporaries (forward and reverse mode both abort on them). This is the allocation-free prerequisite for an exact autodiff Hessian. Pure refactor, identical arithmetic. Verified bit-for-bit: vmec_standalone MHD energy unchanged on solovev (2.548352e+00) and cth_like_fixed_bdy (5.057191e-02). * examples: globalize HVP Newton with a backtracking line search The full Newton step overshoots on stiff 3D equilibria (cth_like stalled at the iteration cap with ||F|| ~ 5e-2). Add a backtracking line search on ||F|| so each step is damped to a decrease. With it the HVP-Newton converges on cth_like in 9 outer iterations (||F|| = 1.8e-10) and still converges solovev in 8. * dft_toroidal: make ForcesToFourier allocation-free The forces transform materialized two per-(surface,m,zeta) Eigen temporaries (tempR_seg, tempZ_seg) inside the inner loop. Reuse per-thread scratch instead, so the whole FFTX-off force path (geometryFromFourier, computeJacobian/Metric/BContra/BCo, pressureAndEnergies, computeMHDForces, forcesToFourier) is now allocation-free end to end. Same arithmetic as the previous .eval(); verified bit-for-bit: solovev 2.548352e+00, cth_like_fixed_bdy 5.057191e-02. * enzyme: exact autodiff of the VMEC Jacobian kernel (forward vs reverse) Demonstrate exact automatic differentiation of a real VMEC nonlinear kernel. JacobianKernel reproduces IdealMhdModel::computeJacobian (half-grid r12/ru12/zu12/rs/zs and the Jacobian tau), written allocation-free over flat buffers, which is the form Enzyme differentiates. For L = 0.5||outputs||^2 the test computes dL/dgeom by reverse mode and the directional derivative dL.v by forward mode, checks both against central finite differences, and against each other: reverse dL.v vs FD : 1.9e-9 forward dL.v vs FD : 1.9e-9 forward vs reverse : 2.9e-15 performance: reverse ~16 us/pass (full gradient), forward ~16 us/pass (one direction) Reverse returns the whole gradient per pass and wins for a scalar gradient; forward is the cheaper primitive for a single Jacobian/Hessian-vector product. tau is nonlinear in the geometry, so this kernel's Jacobian is a genuine building block of the exact MHD force Hessian; the remaining force chain follows the same allocation-free pattern. * ideal_mhd_model: share the Jacobian kernel between solver and autodiff Move the half-grid Jacobian arithmetic into jacobian_kernel.h (ComputeHalfGridJacobian), allocation-free over flat buffers. Production computeJacobian now calls it (followed by the unchanged Jacobian-sign check), and the Enzyme forward/reverse test differentiates the same kernel: one implementation, no duplication. Bit-exact: vmec_standalone MHD energy unchanged on solovev (2.548352e+00) and cth_like_fixed_bdy (5.057191e-02). Autodiff test still matches finite differences and agrees forward vs reverse to 3e-15. * ideal_mhd_model: share the metric kernel (gsqrt, guu, guv, gvv) Extract computeMetricElements into the shared, allocation-free kernel ComputeMetricElements (metric_kernel.h), over flat buffers, and call it from the solver. guv and the 3D part of gvv are computed only when lthreed, matching the original. This is the second force-chain kernel made Enzyme-differentiable (composed into the exact Hessian-vector product later), following the Jacobian kernel pattern. Bit-exact: vmec_standalone MHD energy unchanged on solovev (2.548352e+00, 2D) and cth_like_fixed_bdy (5.057191e-02, 3D path with guv/gvv). * ideal_mhd_model: share the contravariant-field kernel (bsupu, bsupv) Factor the bsupu/bsupv arithmetic out of computeBContra into the shared, allocation-free kernel ComputeBsupContra (bcontra_kernel.h). The lambda normalization (lamscale, + phi') and the chi'/iota profile and toroidal-current-constraint logic stay in the solver verbatim, since they mutate state and update profiles; only the differentiable field arithmetic moves to the shared kernel. Bit-exact across 1 and 4 threads (so the ghost-cell radial partitioning is exercised) on solovev (2.548352e+00, 2D) and cth_like_fixed_bdy (5.057191e-02, 3D). * ideal_mhd_model: share the covariant-field kernel (bsubu, bsubv) Extract the metric index-lowering (bsubu = guu B^u + guv B^v, bsubv = guv B^u + gvv B^v; guv absent in 2D) from computeBCo into the shared, allocation-free kernel ComputeBCo (bco_kernel.h). Bit-exact across 1 and 4 threads on solovev (2.548352e+00) and cth_like_fixed_bdy (5.057191e-02). * ideal_mhd_model: share the magnetic-pressure kernel Extract the field-dependent magnetic pressure |B|^2/2 = 0.5(B^u B_u + B^v B_v) from pressureAndEnergies into the shared, allocation-free kernel ComputeMagneticPressure (pressure_kernel.h). The kinetic-pressure profile and the energy volume integrals stay in the solver. Bit-exact across 1 and 4 threads on solovev (2.548352e+00) and cth_like_fixed_bdy (5.057191e-02). Completes the point-local nonlinear force-chain kernels (Jacobian, metric, B^contra, B_cov, pressure). * ideal_mhd_model: share the MHD force-density kernel Extract computeMHDForces' real-space force-density assembly (armn/azmn/ brmn/bzmn, and crmn/czmn in 3D, even+odd) into the shared, allocation-free kernel ComputeMHDForceDensity (mhdforce_kernel.h). The Eigen arithmetic is preserved verbatim over flat-buffer Eigen::Map views with caller-owned handover/average scratch, so it is bit-for-bit identical. This is the sixth and final point-local force-chain kernel; the six (Jacobian, metric, B^contra, B_cov, pressure, force) now form the local map geometry -> force density, ready to compose into the exact Hessian-vector product. (This branch also merges the allocation-free force kernel, #12, which removes the per-surface heap temporaries this extraction relies on.) Bit-exact across 1 and 4 threads on solovev (2.548352e+00) and cth_like_fixed_bdy (5.057191e-02). * enzyme: exact Hessian of the composed local force map Compose the six shared force-chain kernels (Jacobian, metric, B^contra, B_cov, magnetic pressure, MHD force density) into the single local map g: real-space geometry -> real-space force density, the nonlinear core of VMEC's force. The full MHD force is T^T . g . T with the linear spectral transforms; the exact force Hessian-vector product is therefore T^T . J_g . T . v, and this provides J_g by autodiff. The new test takes the Jacobian of g by forward and reverse Enzyme modes over flat allocation-free buffers, checks both against central finite differences and against each other, and times one forward Jacobian-vector pass against the two force evaluations a finite-difference HVP costs. * ideal_mhd_model: share the hybrid lambda-force kernel Extract hybridLambdaForce's full-grid lambda force (blmn, and clmn in 3D) into lambda_force_kernel.h (ComputeHybridLambdaForce), shared between the solver and the Enzyme autodiff path. The method drops from 115 lines to a single kernel call; the OpenMP barriers stay in the method. The kernel is allocation-free over flat buffers and preserves the radial sweep that carries the inside half-grid point in scratch and shifts it outward each surface, plus the blend of the two bsubv interpolations. This is the lambda-force piece of the augmented functional, the second nonlinear force-density term after the MHD force chain. * ideal_mhd_model: share the constraint-force kernels Extract the two local (non-transform) pieces of the spectral-condensation constraint force into constraint_force_kernel.h, shared between the solver and the Enzyme autodiff path: - ComputeEffectiveConstraintForce: gConEff = (rCon-rCon0) ru + (zCon-zCon0) zu (effectiveConstraintForce), skipping the axis surface. - AddConstraintForces: add the bandpass-filtered gCon back into the MHD R/Z forces and write frcon/fzcon (the constraint part of assembleTotalForces). The Fourier-space bandpass between them stays the shared free function deAliasConstraintForce; the free-boundary rBSq contribution stays in assembleTotalForces. Allocation-free over flat buffers. This completes the local force-density terms of the augmented functional (MHD + lambda + constraint), the nonlinear core of the exact Hessian. * enzyme: extend the composed-force Hessian test with the lambda force Add the hybrid lambda force (lambda_force_kernel.h) to the composed local map g and differentiate the combined MHD-plus-lambda force density by forward and reverse Enzyme modes. This proves J_g for the second nonlinear force-density term, not just the MHD force chain. The spectral-condensation constraint force also carries a linear Fourier bandpass; it is validated end-to-end against the finite-difference HVP in the pybind exact-HVP path rather than in this flat-buffer microtest. * apply pre-commit formatting (ruff, docformatter, clang-format) * apply pre-commit formatting (ruff, docformatter, clang-format) * apply pre-commit formatting (ruff, docformatter, clang-format) * apply pre-commit formatting (ruff, docformatter, clang-format) * apply pre-commit formatting (ruff, docformatter, clang-format) * apply pre-commit formatting (ruff, docformatter, clang-format) * apply pre-commit formatting (ruff, docformatter, clang-format) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * bazel: declare force-chain kernel headers in ideal_mhd_model (sandbox fix) * test: docformatter-format test_internal_gradient docstrings Satisfies the docformatter pre-commit hook (was failing CI). * test: docformatter-format external/internal optimizer test docstrings Satisfies the docformatter pre-commit hook (was failing CI). * ci: re-trigger (transient apt-403 on packages.microsoft.com) * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * ci: skip benchmark result upload on fork PRs (token is read-only) The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing. * ci: build VMEC2000 from source so the compat test runs on numpy 2 The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure. * test: skip vmecpp-only indata fields in the VMEC2000 compat subset With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one. * build: pin abseil to the 20260107.1 commit hash Pin the FetchContent abseil dependency to commit 255c84d (the exact commit behind the 20260107.1 LTS tag) instead of the tag itself, so a moved tag cannot change the dependency under us. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash for Clang >= 21. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash for Clang >= 21. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash for Clang >= 21. * ci: sync VMEC2000-from-source build, benchmark fork guard, abseil commit pin Bring this stack branch up to the corrected CI baseline (from #583/#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash for Clang >= 21. * ci: cache and pin the VMEC2000-from-source build Use the canonical recipe (cache the built wheel keyed on the pinned source commit 728af8b, drop the unused FFTW/HDF5 dev packages) instead of rebuilding VMEC2000 unpinned on every run. * ideal_mhd_model: mark Jacobian kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: hoist ForcesToFourier scratch out of the inner loop The allocation-free rewrite placed tempR_seg/tempZ_seg in a block-scope thread_local inside the (jF, m, zeta) inner loop, which emits a __tls_get_addr call and an init-guard branch every iteration. Declare the two scratch vectors once at function scope instead: still allocation-free in the hot loop and per-thread safe via the stack frame, without the per-iteration TLS overhead. Same arithmetic; cma and w7x wout are bit-for-bit unchanged. * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: hoist ForcesToFourier scratch out of the inner loop The allocation-free rewrite placed tempR_seg/tempZ_seg in a block-scope thread_local inside the (jF, m, zeta) inner loop, which emits a __tls_get_addr call and an init-guard branch every iteration. Declare the two scratch vectors once at function scope instead: still allocation-free in the hot loop and per-thread safe via the stack frame, without the per-iteration TLS overhead. Same arithmetic; cma and w7x wout are bit-for-bit unchanged. * ideal_mhd_model: hoist ForcesToFourier scratch out of the inner loop The allocation-free rewrite placed tempR_seg/tempZ_seg in a block-scope thread_local inside the (jF, m, zeta) inner loop, which emits a __tls_get_addr call and an init-guard branch every iteration. Declare the two scratch vectors once at function scope instead: still allocation-free in the hot loop and per-thread safe via the stack frame, without the per-iteration TLS overhead. Same arithmetic; cma and w7x wout are bit-for-bit unchanged. * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: hoist ForcesToFourier scratch out of the inner loop The allocation-free rewrite placed tempR_seg/tempZ_seg in a block-scope thread_local inside the (jF, m, zeta) inner loop, which emits a __tls_get_addr call and an init-guard branch every iteration. Declare the two scratch vectors once at function scope instead: still allocation-free in the hot loop and per-thread safe via the stack frame, without the per-iteration TLS overhead. Same arithmetic; cma and w7x wout are bit-for-bit unchanged. * ideal_mhd_model: hoist ForcesToFourier scratch out of the inner loop The allocation-free rewrite placed tempR_seg/tempZ_seg in a block-scope thread_local inside the (jF, m, zeta) inner loop, which emits a __tls_get_addr call and an init-guard branch every iteration. Declare the two scratch vectors once at function scope instead: still allocation-free in the hot loop and per-thread safe via the stack frame, without the per-iteration TLS overhead. Same arithmetic; cma and w7x wout are bit-for-bit unchanged. * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * ideal_mhd_model: mark Jacobian metric kernel buffers __restrict Raw double* kernel params over the same flat layout prevent the compiler from vectorizing the pointwise loop (assumed aliasing), so on w7x these kernels ran ~2x slower than the Eigen-expression code they replaced. The buffers never overlap; mark them __restrict to restore SIMD. Enzyme derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark). * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * enzyme: run the AD smoke test through bazel instead of ctest Move the Enzyme autodiff smoke test into the bazel test framework, which owns every other C++ test in this repository, and drop the separate CMake ctest path that nothing in CI exercised. - vmecpp/common/enzyme/BUILD.bazel: an `enzyme` header library plus an `enzyme_smoke_test` cc_test. The test is tagged `manual` so the default GCC `bazel test //...` skips it (the Enzyme intrinsics only resolve under Clang with the plugin attached) and never tries to compile it with GCC. - .bazelrc: a `--config=enzyme` that sets -O2 so the Enzyme optimization pass fires. Select Clang with CC/CXX and pass the plugin path the way -DVMECPP_ENZYME_PLUGIN did under CMake: CC=clang CXX=clang++ bazel test --config=enzyme \ --copt=-fplugin=/path/to/ClangEnzyme-NN.so \ //vmecpp/common/enzyme:enzyme_smoke_test - CMakeLists.txt: remove the VMECPP_ENABLE_ENZYME option and the ctest registration it only existed to drive. * ci: build ClangEnzyme and run the enzyme smoke test in CI Add a GitHub Actions job that gives the Enzyme autodiff smoke test actual CI coverage. It mirrors the EnzymeAD upstream recipe: install Clang/LLVM 21 from apt.llvm.org, build a pinned ClangEnzyme-21 plugin (v0.0.264, the version this stack is developed against) against the installed LLVM and Clang, then run the bazel target under --config=enzyme with the plugin attached. The plugin build is cached on the pinned ref so only the first run pays for it. This is what the enzyme test needed beyond the bazel move: the default GCC test_bazel job skips the manual-tagged target, so without a Clang/Enzyme job nothing exercised it. * output_quantities: compare jcuru/jcurv at the standard tolerance The Jacobian-kernel refactor is structure-only, so drop the opt-in current_density_tolerance loosening and compare current densities at the same relabs tolerance as every other wout quantity. * test: address review nits in test_internal_gradient Drop the local-dev ImportError fallback (use the canonical vmecpp.cpp import as elsewhere) and the redundant __main__ block, and note that the raw and preconditioned forces both vanish at convergence. * enzyme: drop timing-dependent benchmark from local force Hessian test Remove the chrono-based forward-JVP vs FD-HVP timing loop. Timing assertions are environment-dependent and unfit as blocking unit tests; the test keeps the forward/reverse/finite-difference correctness checks. Per-machine cost numbers belong in the non-blocking benchmark harness. * enzyme: drop timing-dependent benchmark from local force Hessian test Remove the chrono-based forward-JVP vs FD-HVP timing loop. Timing assertions are environment-dependent and unfit as blocking unit tests; the test keeps the forward/reverse/finite-difference correctness checks. Per-machine cost numbers belong in the non-blocking benchmark harness. * enzyme: drop timing-dependent benchmark from local force Hessian test Remove the chrono-based forward-JVP vs FD-HVP timing loop. Timing assertions are environment-dependent and unfit as blocking unit tests; the test keeps the forward/reverse/finite-difference correctness checks. Per-machine cost numbers belong in the non-blocking benchmark harness. * enzyme: drop timing-dependent benchmark from local force Hessian test Remove the chrono-based forward-JVP vs FD-HVP timing loop. Timing assertions are environment-dependent and unfit as blocking unit tests; the test keeps the forward/reverse/finite-difference correctness checks. Per-machine cost numbers belong in the non-blocking benchmark harness. * ideal_mhd_model: drop shadowed dead scratch in toroidal force DFT The function-scope tempR_seg/tempZ_seg were never read: the inner loop declares its own thread_local scratch of the same name that shadows them. Remove the unused pair and its inaccurate comment; the thread_local scratch in the inner loop is the one actually reused across iterations. * ideal_mhd_model: drop shadowed dead scratch in toroidal force DFT The function-scope tempR_seg/tempZ_seg were never read: the inner loop declares its own thread_local scratch of the same name that shadows them. Remove the unused pair and its inaccurate comment; the thread_local scratch in the inner loop is the one actually reused across iterations. * ideal_mhd_model: drop shadowed dead scratch in toroidal force DFT The function-scope tempR_seg/tempZ_seg were never read: the inner loop declares its own thread_local scratch of the same name that shadows them. Remove the unused pair and its inaccurate comment; the thread_local scratch in the inner loop is the one actually reused across iterations. * ideal_mhd_model: drop shadowed dead scratch in toroidal force DFT The function-scope tempR_seg/tempZ_seg were never read: the inner loop declares its own thread_local scratch of the same name that shadows them. Remove the unused pair and its inaccurate comment; the thread_local scratch in the inner loop is the one actually reused across iterations. * ideal_mhd_model: drop shadowed dead scratch in toroidal force DFT The function-scope tempR_seg/tempZ_seg were never read: the inner loop declares its own thread_local scratch of the same name that shadows them. Remove the unused pair and its inaccurate comment; the thread_local scratch in the inner loop is the one actually reused across iterations. * examples: cold-start external optimizers from axis-less inputs make_model now reguesses the magnetic axis when the initial geometry has a singular Jacobian, mirroring the native solver's first-iterate axis reguess (vmec.cc SolveEquilibriumLoop). Inputs that ship no axis (raxis/zaxis all zero, e.g. cma.json) otherwise return a zero raw force at the BAD_JACOBIAN checkpoint. Add a cma test (3D stellarator, nfp=2, ntor=6) that exercises the non-axisymmetric force chain: after the reguess the raw internal-basis force and the Hessian-vector product are finite and nonzero. * examples: cold-start external optimizers from axis-less inputs make_model now reguesses the magnetic axis when the initial geometry has a singular Jacobian, mirroring the native solver's first-iterate axis reguess (vmec.cc SolveEquilibriumLoop). Inputs that ship no axis (raxis/zaxis all zero, e.g. cma.json) otherwise return a zero raw force at the BAD_JACOBIAN checkpoint. Add a cma test (3D stellarator, nfp=2, ntor=6) that exercises the non-axisymmetric force chain: after the reguess the raw internal-basis force and the Hessian-vector product are finite and nonzero. * examples: cold-start external optimizers from axis-less inputs make_model now reguesses the magnetic axis when the initial geometry has a singular Jacobian, mirroring the native solver's first-iterate axis reguess (vmec.cc SolveEquilibriumLoop). Inputs that ship no axis (raxis/zaxis all zero, e.g. cma.json) otherwise return a zero raw force at the BAD_JACOBIAN checkpoint. Add a cma test (3D stellarator, nfp=2, ntor=6) that exercises the non-axisymmetric force chain: after the reguess the raw internal-basis force and the Hessian-vector product are finite and nonzero. * examples: docformatter-compliant docstrings (pinned pre-commit hook) * examples: docformatter-compliant docstrings (pinned pre-commit hook) * examples: docformatter-compliant docstrings (pinned pre-commit hook) * ci: re-trigger checks (Pre-commit run was stale; tree is docformatter-clean) * ci: re-trigger asan (vmec_in_memory_mgrid_test jcuru was at the 1e-7 boundary) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * output_quantities: compare jcuru/jcurv at a looser opt-in tolerance The free-boundary in-memory-vs-disk mgrid golden compares two independent solves. jcuru/jcurv are curl(B) current densities that amplify the rounding of the converged state, so under vectorized/optimized builds the two paths diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other wout quantity still agrees to 1e-7. The math is unchanged: with vs without the kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so this is an FP-ordering reproducibility floor, not an accuracy regression. Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the main tolerance, so every other caller is unchanged) and have the two vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping 1e-7 for all profiles and geometry. (cherry picked from commit 27d36d21e1dd8ea6f73127b95bdc81d529f81672) * ideal_mhd_model: include contravariant kernel header --------- Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com> --- examples/external_optimizers.py | 244 ++++++++++++++++++ .../cpp/vmecpp/vmec/pybind11/pybind_vmec.cc | 63 +++++ tests/test_external_optimizers.py | 98 +++++++ tests/test_hessian.py | 64 +++++ tests/test_preconditioner.py | 63 +++++ 5 files changed, 532 insertions(+) create mode 100644 examples/external_optimizers.py create mode 100644 tests/test_external_optimizers.py create mode 100644 tests/test_hessian.py create mode 100644 tests/test_preconditioner.py diff --git a/examples/external_optimizers.py b/examples/external_optimizers.py new file mode 100644 index 000000000..4e06bb2ea --- /dev/null +++ b/examples/external_optimizers.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# 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. +""" + +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 + +from vmecpp.cpp import _vmecpp # type: ignore[import] + +DEFAULT_INPUT = ( + Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.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_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=300, 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 and is + # state-invariant once assembled. + 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_preconditioned_descent, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + 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() diff --git a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc index 063ac9ca0..cda5d0ad9 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc @@ -227,6 +227,7 @@ class VmecModel { // optimizers minimizing the MHD energy functional need; mhd_energy is already // set earlier in update(), so it is valid at the checkpoint too. void Evaluate(int iter1, int iter2, bool precondition = true) { + ++force_eval_count_; bool need_restart = false; std::string error_message; const vmecpp::VmecCheckpoint checkpoint = @@ -266,6 +267,12 @@ class VmecModel { } bool need_restart() const { return last_need_restart_; } + // Total forward-model (force) evaluations since construction or the last + // reset. Counts every Evaluate, including those inside hessian_vector_product + // and preconditioner assembly, for a fair cross-optimizer cost comparison. + long force_eval_count() const { return force_eval_count_; } + void reset_force_eval_count() { force_eval_count_ = 0; } + // The Garabedian-style time step (PerformTimeStep): for each Fourier // coefficient, v = velocity_scale*(conjugation*v + dt*force); x += dt*v. void PerformTimeStep(double velocity_scale, double conjugation_parameter, @@ -394,6 +401,55 @@ class VmecModel { return FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_); } + // Hessian-vector product of VMEC's augmented functional, computed inside + // VMEC++ by a central directional derivative of the analytic force (which is + // the gradient): H v = (F(x + eps v) - F(x - eps v)) / (2 eps), in the + // decomposed internal basis. This is the matrix-free Hessian information an + // internal or external Newton-Krylov solver needs; F itself is exact, so only + // the directional step is finite-differenced. The current state is restored. + Eigen::VectorXd HessianVectorProduct(const Eigen::VectorXd &v, + double eps_rel = 1e-7) { + const Eigen::VectorXd x = + FlattenActive(*vmec_->decomposed_x_[0], vmec_->s_); + const double vnorm = v.norm(); + if (vnorm == 0.0) { + return Eigen::VectorXd::Zero(x.size()); + } + const double eps = eps_rel * (1.0 + x.norm()) / vnorm; + UnflattenActive(*vmec_->decomposed_x_[0], vmec_->s_, x + eps * v); + Evaluate(2, 2, /*precondition=*/false); + const Eigen::VectorXd fp = + FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_); + UnflattenActive(*vmec_->decomposed_x_[0], vmec_->s_, x - eps * v); + Evaluate(2, 2, /*precondition=*/false); + const Eigen::VectorXd fm = + FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_); + UnflattenActive(*vmec_->decomposed_x_[0], vmec_->s_, x); + return (fp - fm) / (2.0 * eps); + } + + // Apply VMEC's preconditioner M^-1 to a vector in the decomposed internal + // basis, mirroring the native apply sequence (m=1, radial, lambda). This is + // VMEC's hand-built approximate inverse Hessian; gradient-based solvers use + // it as the metric (preconditioned Krylov / quasi-Newton, and as the + // preconditioner for the Hessian solve in adjoint sensitivities). + // + // Requires a prior evaluate(precondition=true) at the current state: the + // radial preconditioner is assembled inside that forward-model call. + Eigen::VectorXd ApplyPreconditioner(const Eigen::VectorXd &v) const { + vmecpp::FourierForces tmp(&vmec_->s_, vmec_->r_[0].get(), vmec_->fc_.ns); + tmp.setZero(); + UnflattenActive(tmp, vmec_->s_, v); + vmecpp::IdealMhdModel &model = *vmec_->m_[0]; + model.applyM1Preconditioner(tmp); + const absl::Status status = model.applyRZPreconditioner(tmp); + if (!status.ok()) { + throw std::runtime_error(std::string(status.message())); + } + model.applyLambdaPreconditioner(tmp); + return FlattenActive(tmp, vmec_->s_); + } + // Residuals (set by Evaluate()): invariant {fsqr,fsqz,fsql} and // preconditioned {fsqr1,fsqz1,fsql1}. double fsqr() const { return vmec_->fc_.fsqr; } @@ -466,6 +522,7 @@ class VmecModel { int last_preconditioner_update_ = 0; int last_full_update_nestor_ = 0; bool last_need_restart_ = false; + long force_eval_count_ = 0; }; } // anonymous namespace @@ -1219,6 +1276,12 @@ PYBIND11_MODULE(_vmecpp, m) { .def("get_state", &VmecModel::GetState) .def("set_state", &VmecModel::SetState, py::arg("state")) .def("get_forces", &VmecModel::GetForces) + .def("apply_preconditioner", &VmecModel::ApplyPreconditioner, + py::arg("v")) + .def("hessian_vector_product", &VmecModel::HessianVectorProduct, + py::arg("v"), py::arg("eps_rel") = 1e-7) + .def_property_readonly("force_eval_count", &VmecModel::force_eval_count) + .def("reset_force_eval_count", &VmecModel::reset_force_eval_count) .def_property_readonly("fsqr", &VmecModel::fsqr) .def_property_readonly("fsqz", &VmecModel::fsqz) .def_property_readonly("fsql", &VmecModel::fsql) diff --git a/tests/test_external_optimizers.py b/tests/test_external_optimizers.py new file mode 100644 index 000000000..8b5112818 --- /dev/null +++ b/tests/test_external_optimizers.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""External optimizers reach the same equilibrium as the native solver. + +The raw internal-basis force (gradient of VMEC's augmented functional) is the residual +F(x); F(x) = 0 at equilibrium. Both a native-style preconditioned descent and a +Jacobian-free Newton-Krylov solver drive it to zero and recover the native solver's +converged state and energy. +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from external_optimizers import ( # type: ignore + make_model, + reference_equilibrium, + residual, + solve_newton_hvp, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_preconditioned_descent, +) + +CMA = ( + Path(__file__).resolve().parents[1] + / "src" + / "vmecpp" + / "cpp" + / "vmecpp" + / "test_data" + / "cma.json" +) + + +@pytest.fixture(scope="module") +def reference(): + return reference_equilibrium() + + +@pytest.mark.parametrize( + "solver", + [ + solve_preconditioned_descent, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_newton_hvp, + ], +) +def test_optimizer_reaches_equilibrium(solver, reference): + x_star, w_star = reference + x, result = solver() + # Force balance achieved. + assert result.residual_norm < 1e-7 + # Same equilibrium as the native solver. + assert abs(result.energy - w_star) < 1e-8 + assert np.linalg.norm(x - x_star) < 1e-5 + + +def test_preconditioner_accelerates_newton_krylov(): + # VMEC's preconditioner is the inverse-Hessian approximation: using it as the + # inner Krylov preconditioner cuts the force evaluations substantially. + _, plain = solve_newton_krylov() + _, precond = solve_newton_krylov_preconditioned() + assert precond.force_evals < plain.force_evals + + +def test_newton_hvp_converges_in_few_iterations(): + # True Newton with VMEC++'s own Hessian-vector product converges in a handful + # of outer iterations (second-order), far fewer than first-order descent. + _, newton = solve_newton_hvp() + _, descent = solve_preconditioned_descent() + assert newton.outer_iters < 20 + assert newton.outer_iters < descent.outer_iters + + +def test_cma_cold_start_exercises_non_axisymmetric_paths(): + # cma.json is a 3D stellarator (nfp=2, ntor=6) that ships no magnetic axis + # (raxis/zaxis all zero), so the initial geometry has a singular Jacobian. + # make_model reguesses the axis like the native solver, after which the raw + # internal-basis force and the Hessian-vector product are well defined on the + # non-axisymmetric force chain. + model = make_model(CMA, ns=25) + x0 = np.asarray(model.get_state(), float) + f0 = residual(model)(x0) + assert np.all(np.isfinite(f0)) + assert np.linalg.norm(f0) > 0.0 + + rng = np.random.default_rng(0) + v = rng.standard_normal(x0.size) + hv = np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float) + assert np.all(np.isfinite(hv)) + assert np.linalg.norm(hv) > 0.0 diff --git a/tests/test_hessian.py b/tests/test_hessian.py new file mode 100644 index 000000000..8034f5ace --- /dev/null +++ b/tests/test_hessian.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VmecModel.hessian_vector_product gives the augmented functional's curvature. + +The Hessian-vector product is a central directional derivative of the analytic +force (the gradient of VMEC's augmented functional), computed inside VMEC++: +H v = (F(x + eps v) - F(x - eps v)) / (2 eps). It must be linear in v and agree +with an independent finite difference of the force, and it restores the state. +""" + +from pathlib import Path + +import numpy as np + +from vmecpp.cpp import _vmecpp # type: ignore + +SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" + + +def _model(ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(SOLOVEV)), ns) + + +def _raw_force(model, x): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), float) + + +def test_hvp_matches_finite_difference(): + m = _model() + m.evaluate(2, 2, False) + x = np.asarray(m.get_state(), float) + rng = np.random.default_rng(0) + v = rng.standard_normal(x.size) + v /= np.linalg.norm(v) + + hv = np.asarray(m.hessian_vector_product(np.ascontiguousarray(v)), float) + + eps = 1e-6 + fd = (_raw_force(m, x + eps * v) - _raw_force(m, x - eps * v)) / (2 * eps) + assert np.linalg.norm(hv - fd) < 1e-5 * np.linalg.norm(fd) + + +def test_hvp_is_linear(): + m = _model() + m.evaluate(2, 2, False) + rng = np.random.default_rng(1) + v = rng.standard_normal(np.asarray(m.get_state()).size) + hv = np.asarray(m.hessian_vector_product(np.ascontiguousarray(v)), float) + hv2 = np.asarray(m.hessian_vector_product(np.ascontiguousarray(2.0 * v)), float) + assert np.linalg.norm(hv2 - 2.0 * hv) < 1e-9 * np.linalg.norm(hv) + + +def test_hvp_restores_state(): + m = _model() + m.evaluate(2, 2, False) + x0 = np.asarray(m.get_state(), float).copy() + rng = np.random.default_rng(2) + v = rng.standard_normal(x0.size) + m.hessian_vector_product(np.ascontiguousarray(v)) + assert np.allclose(np.asarray(m.get_state(), float), x0) diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py new file mode 100644 index 000000000..7009dee8a --- /dev/null +++ b/tests/test_preconditioner.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VmecModel.apply_preconditioner exposes VMEC's preconditioner as an operator. + +The preconditioner M^-1 is VMEC's hand-built approximate inverse Hessian. The native +solver applies it to the raw force to get its search direction, so +apply_preconditioner(raw force) must equal the preconditioned force exactly. The +operator is linear and, once assembled (via evaluate(precondition=True)), does not +depend on the current state, so it can be reused as a frozen preconditioner for +Krylov/quasi-Newton solvers. +""" + +from pathlib import Path + +import numpy as np + +try: + from vmecpp.cpp import _vmecpp +except ImportError: + import _vmecpp + +SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" + + +def _model(ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(SOLOVEV)), ns) + + +def test_preconditioner_matches_native_search_direction(): + m = _model() + m.evaluate(2, 2, True) # assemble preconditioner + preconditioned force + f_prec = np.asarray(m.get_forces(), float) + m.evaluate(2, 2, False) # raw force (does not reassemble) + f_raw = np.asarray(m.get_forces(), float) + minv_fraw = np.asarray(m.apply_preconditioner(f_raw), float) + assert np.linalg.norm(minv_fraw - f_prec) <= 1e-12 * np.linalg.norm(f_prec) + + +def test_preconditioner_is_linear_and_finite(): + m = _model() + m.evaluate(2, 2, True) + rng = np.random.default_rng(0) + v = np.ascontiguousarray(rng.standard_normal(np.asarray(m.get_state()).size)) + mv = np.asarray(m.apply_preconditioner(v), float) + m2v = np.asarray(m.apply_preconditioner(np.ascontiguousarray(2.0 * v)), float) + assert np.all(np.isfinite(mv)) + assert np.linalg.norm(m2v - 2.0 * mv) <= 1e-12 * np.linalg.norm(mv) + + +def test_preconditioner_state_invariant_after_assembly(): + m = _model() + m.evaluate(2, 2, True) + rng = np.random.default_rng(1) + x = np.asarray(m.get_state(), float) + v = np.ascontiguousarray(rng.standard_normal(x.size)) + mv0 = np.asarray(m.apply_preconditioner(v), float) + # Move to a different state and raw-evaluate (no reassembly). + m.set_state(np.ascontiguousarray(x + 0.01 * rng.standard_normal(x.size))) + m.evaluate(2, 2, False) + mv1 = np.asarray(m.apply_preconditioner(v), float) + assert np.linalg.norm(mv1 - mv0) <= 1e-12 * np.linalg.norm(mv0) From b05e2daa3fc83aefc301d74237fa82095250933c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Jura=C5=A1i=C4=87?= <166746189+jurasic-pf@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:46:02 +0200 Subject: [PATCH 02/12] 3d external optimizer example (#618) --- examples/external_optimizers.py | 49 +++++++++++++++++++++++++++++---- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/examples/external_optimizers.py b/examples/external_optimizers.py index 4e06bb2ea..b4da85f26 100644 --- a/examples/external_optimizers.py +++ b/examples/external_optimizers.py @@ -24,6 +24,15 @@ 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 @@ -36,10 +45,14 @@ 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" / "solovev.json" + Path(__file__).resolve().parents[1] + / "examples" + / "data" + / "cth_like_fixed_bdy.json" ) @@ -111,6 +124,29 @@ def _finish(model, name, x, outer_iters, t0): ) +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 ): @@ -134,7 +170,7 @@ def solve_preconditioned_descent( def solve_newton_krylov( - input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=300, preconditioned=False + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=2500, preconditioned=False ): model = make_model(input_path, ns) F = residual(model) @@ -143,8 +179,10 @@ def solve_newton_krylov( 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 and is - # state-invariant once assembled. + # 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] @@ -218,9 +256,10 @@ def solve_newton_hvp( ALL_SOLVERS = ( + solve_vmecpp, solve_preconditioned_descent, - solve_newton_krylov, solve_newton_krylov_preconditioned, + solve_newton_krylov, solve_newton_hvp, ) From 4a9e717ae94964d78908d9db54d5c8ff75ea01f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Jura=C5=A1i=C4=87?= <166746189+jurasic-pf@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:49:43 +0200 Subject: [PATCH 03/12] Remove explicit test of an example script (#622) --- tests/test_external_optimizers.py | 98 ------------------------------- 1 file changed, 98 deletions(-) delete mode 100644 tests/test_external_optimizers.py diff --git a/tests/test_external_optimizers.py b/tests/test_external_optimizers.py deleted file mode 100644 index 8b5112818..000000000 --- a/tests/test_external_optimizers.py +++ /dev/null @@ -1,98 +0,0 @@ -# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH -# -# -# SPDX-License-Identifier: MIT -"""External optimizers reach the same equilibrium as the native solver. - -The raw internal-basis force (gradient of VMEC's augmented functional) is the residual -F(x); F(x) = 0 at equilibrium. Both a native-style preconditioned descent and a -Jacobian-free Newton-Krylov solver drive it to zero and recover the native solver's -converged state and energy. -""" - -import sys -from pathlib import Path - -import numpy as np -import pytest - -sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) -from external_optimizers import ( # type: ignore - make_model, - reference_equilibrium, - residual, - solve_newton_hvp, - solve_newton_krylov, - solve_newton_krylov_preconditioned, - solve_preconditioned_descent, -) - -CMA = ( - Path(__file__).resolve().parents[1] - / "src" - / "vmecpp" - / "cpp" - / "vmecpp" - / "test_data" - / "cma.json" -) - - -@pytest.fixture(scope="module") -def reference(): - return reference_equilibrium() - - -@pytest.mark.parametrize( - "solver", - [ - solve_preconditioned_descent, - solve_newton_krylov, - solve_newton_krylov_preconditioned, - solve_newton_hvp, - ], -) -def test_optimizer_reaches_equilibrium(solver, reference): - x_star, w_star = reference - x, result = solver() - # Force balance achieved. - assert result.residual_norm < 1e-7 - # Same equilibrium as the native solver. - assert abs(result.energy - w_star) < 1e-8 - assert np.linalg.norm(x - x_star) < 1e-5 - - -def test_preconditioner_accelerates_newton_krylov(): - # VMEC's preconditioner is the inverse-Hessian approximation: using it as the - # inner Krylov preconditioner cuts the force evaluations substantially. - _, plain = solve_newton_krylov() - _, precond = solve_newton_krylov_preconditioned() - assert precond.force_evals < plain.force_evals - - -def test_newton_hvp_converges_in_few_iterations(): - # True Newton with VMEC++'s own Hessian-vector product converges in a handful - # of outer iterations (second-order), far fewer than first-order descent. - _, newton = solve_newton_hvp() - _, descent = solve_preconditioned_descent() - assert newton.outer_iters < 20 - assert newton.outer_iters < descent.outer_iters - - -def test_cma_cold_start_exercises_non_axisymmetric_paths(): - # cma.json is a 3D stellarator (nfp=2, ntor=6) that ships no magnetic axis - # (raxis/zaxis all zero), so the initial geometry has a singular Jacobian. - # make_model reguesses the axis like the native solver, after which the raw - # internal-basis force and the Hessian-vector product are well defined on the - # non-axisymmetric force chain. - model = make_model(CMA, ns=25) - x0 = np.asarray(model.get_state(), float) - f0 = residual(model)(x0) - assert np.all(np.isfinite(f0)) - assert np.linalg.norm(f0) > 0.0 - - rng = np.random.default_rng(0) - v = rng.standard_normal(x0.size) - hv = np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float) - assert np.all(np.isfinite(hv)) - assert np.linalg.norm(hv) > 0.0 From 2e4c22dec2a1ca5617ba5d4a9abd6f549bce4abf Mon Sep 17 00:00:00 2001 From: clazzati-pf Date: Thu, 9 Jul 2026 18:13:17 +0200 Subject: [PATCH 04/12] Refactor Fourier-resolution continuation into vmecpp.run() (#619) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit moved vmecpp.run_continuation() logic directly into vmecpp.run() Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com> --- examples/fourier_resolution_increase.py | 17 ++- src/vmecpp/__init__.py | 78 +++++++++++-- src/vmecpp/_continuation.py | 101 +++++++++-------- src/vmecpp/simsopt_compat.py | 27 +++-- tests/test_continuation.py | 145 +++++++++++++++++------- 5 files changed, 250 insertions(+), 118 deletions(-) diff --git a/examples/fourier_resolution_increase.py b/examples/fourier_resolution_increase.py index 69dda7c91..375452153 100644 --- a/examples/fourier_resolution_increase.py +++ b/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)) diff --git a/src/vmecpp/__init__.py b/src/vmecpp/__init__.py index 9883157d5..492b87cc9 100644 --- a/src/vmecpp/__init__.py +++ b/src/vmecpp/__init__.py @@ -20,7 +20,7 @@ import pydantic from vmecpp import _util -from vmecpp._continuation import interpolate_solution, run_continuation +from vmecpp._continuation import _run_fourier_continuation, interpolate_solution from vmecpp._free_boundary import ( MagneticFieldResponseTable, MakegridParameters, @@ -77,6 +77,32 @@ def _wrap_int_as_float( pydantic.BeforeValidator(lambda x: np.array(x).astype(np.int64)), ] + +def _coerce_mpol_ntor(value: typing.Any) -> int | np.ndarray: + """Normalizes an ``mpol``/``ntor`` field value. + + A length-1 sequence is equivalent to a scalar and is collapsed to one; anything + longer is kept as an int array representing a per-``ns_array``-step Fourier + resolution continuation schedule. + """ + if isinstance(value, (int, np.integer)): + return int(value) + array = np.atleast_1d(np.asarray(value, dtype=np.int64)) + return int(array[0]) if array.size == 1 else array + + +MpolNtorField: typing.TypeAlias = typing.Annotated[ + int | jt.Int[np.ndarray, "num_fourier_steps"], + pydantic.BeforeValidator(_coerce_mpol_ntor), +] + + +def _final_resolution(value: int | np.ndarray) -> int: + """The target Fourier resolution: itself if scalar, else the schedule's last + (finest) entry.""" + return value if isinstance(value, int) else int(value[-1]) + + AuxFType = typing.Annotated[ _ArrayType, pydantic.BeforeValidator(lambda x: _util.right_pad(x, ndfmax, 0.0)), @@ -188,12 +214,23 @@ class VmecInput(BaseModelWithNumpy): nfp: int = 1 """Number of toroidal field periods (=1 for Tokamak)""" - mpol: int = 6 - """Number of poloidal Fourier harmonics; m = 0, 1, ..., (mpol-1)""" + mpol: MpolNtorField = 6 + """Number of poloidal Fourier harmonics; m = 0, 1, ..., (mpol-1). - ntor: int = 0 + May also be a sequence of ints, with one entry per ``ns_array`` step (a scalar + broadcasts to every step), to request continuation in Fourier resolution: + ``vmecpp.run()`` then solves each step in turn, hot-restarting from the + previous step's solution interpolated to the new resolution (see + :func:`interpolate_solution`). The boundary coefficients (``rbc``, ``zbs``, ...) + are always defined at the final (largest-index) entry's resolution. + """ + + ntor: MpolNtorField = 0 """Number of toroidal Fourier harmonics; n = -ntor, -ntor+1, ..., -1, 0, 1, ..., - ntor-1, ntor.""" + ntor-1, ntor. + + May be a sequence of ints, analogous to :attr:`mpol`; see its docstring. + """ mpol_geometry: int = -1 """Optional reduced poloidal resolution for the geometry (R, Z). @@ -465,7 +502,9 @@ def _validate_fourier_coefficients_shapes(self) -> VmecInput: if self.lasym: mpol_two_ntor_plus_one_fields.extend(["rbs", "zbc"]) - expected_shape = (self.mpol, 2 * self.ntor + 1) + mpol_final = _final_resolution(self.mpol) + ntor_final = _final_resolution(self.ntor) + expected_shape = (mpol_final, 2 * ntor_final + 1) for field in mpol_two_ntor_plus_one_fields: current_value = getattr(self, field) @@ -480,8 +519,8 @@ def _validate_fourier_coefficients_shapes(self) -> VmecInput: field, VmecInput.resize_2d_coeff( current_value, - mpol_new=self.mpol, - ntor_new=self.ntor, + mpol_new=mpol_final, + ntor_new=ntor_final, ), ) return self @@ -634,7 +673,9 @@ def _to_cpp_vmecindata(self) -> _vmecpp.VmecINDATA: ) # this also resizes the readonly_attrs - cpp_indata._set_mpol_ntor(self.mpol, self.ntor) + cpp_indata._set_mpol_ntor( + _final_resolution(self.mpol), _final_resolution(self.ntor) + ) for attr in readonly_attrs - {"mpol", "ntor"}: # now we can set the elements of the readonly_attrs value = getattr(self, attr) @@ -2234,6 +2275,14 @@ def run( restart_from: if present, VMEC++ is initialized using the converged equilibrium from the provided VmecOutput. This can dramatically decrease the number of iterations to 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. + + 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` + entry (a scalar mpol/ntor broadcasts to every step), and each step is solved in turn, + hot-restarting from the previous step's solution interpolated to the new resolution (see + `interpolate_solution`). Example: >>> import vmecpp @@ -2244,6 +2293,16 @@ def run( 0.2033313711 """ input = VmecInput.model_validate(input) + + if not isinstance(input.mpol, int) or not isinstance(input.ntor, int): + return _run_fourier_continuation( + input, + magnetic_field, + max_threads=max_threads, + verbose=verbose, + restart_from=restart_from, + ) + cpp_indata = input._to_cpp_vmecindata() if restart_from is None: @@ -2490,7 +2549,6 @@ def set_profile( # items in the generated documentation. __all__ = [ # noqa: RUF022 "run", - "run_continuation", "interpolate_solution", "VmecInput", "VmecOutput", diff --git a/src/vmecpp/_continuation.py b/src/vmecpp/_continuation.py index b80cb5539..a4fe57d1f 100644 --- a/src/vmecpp/_continuation.py +++ b/src/vmecpp/_continuation.py @@ -1,10 +1,16 @@ """Generalized resolution interpolation and the Python-side continuation driver. VMEC++ converges much more reliably when a hard equilibrium is approached through -a sequence of increasing resolutions (the classic ``ns_array`` multi-grid, and now -also ``mpol_array`` / ``ntor_array`` Fourier continuation). Each step solves a single -resolution and hot-restarts from the previous step's solution, interpolated to the -new resolution by :func:`interpolate_solution`. +a sequence of increasing resolutions (the classic ``ns_array`` multi-grid, and also +Fourier continuation via a sequence-valued ``VmecInput.mpol`` / ``.ntor``). Each step +solves a single resolution and hot-restarts from the previous step's solution, +interpolated to the new resolution by :func:`interpolate_solution`. + +:func:`vmecpp.run` dispatches to :func:`_run_fourier_continuation` (this module) +whenever ``input.mpol`` and/or ``input.ntor`` is a sequence rather than a plain int; +:func:`interpolate_solution` itself is public API and can also be used directly, e.g. +to hand-roll a custom continuation schedule (see +``examples/fourier_resolution_increase.py``). The interpolation is purely a Python operation on a converged :class:`VmecOutput`: the flux-surface geometry is interpolated radially along the normalized toroidal @@ -22,7 +28,8 @@ import numpy as np if typing.TYPE_CHECKING: - from vmecpp import VmecInput, VmecOutput + from vmecpp import OutputMode, VmecInput, VmecOutput + from vmecpp._free_boundary import MagneticFieldResponseTable # State-vector geometry arrays, shape [mn_mode, n_surfaces]. These are the only # quantities VMEC++ reads back when hot-restarting, so they must be interpolated. @@ -285,63 +292,59 @@ def _step_input( return step -def run_continuation( +def _run_fourier_continuation( input: VmecInput, + magnetic_field: MagneticFieldResponseTable | None, *, - ns_array: typing.Sequence[int] | None = None, - mpol_array: typing.Sequence[int] | None = None, - ntor_array: typing.Sequence[int] | None = None, - ftol_array: typing.Sequence[float] | None = None, - niter_array: typing.Sequence[int] | None = None, - **run_kwargs: typing.Any, + max_threads: int | None, + verbose: bool | int | OutputMode, + restart_from: VmecOutput | None, ) -> VmecOutput: - """Solve an equilibrium by continuation in radial and Fourier resolution. + """Solves an equilibrium by continuation in Fourier resolution. + Called by :func:`vmecpp.run` whenever ``input.mpol`` and/or ``input.ntor`` is a + sequence rather than a plain int. Each entry pairs with the corresponding + ``input.ns_array`` entry (a scalar ``mpol``/``ntor`` broadcasts to every step). Each step solves a single ``(ns, mpol, ntor)`` resolution and hot-restarts from the previous step's solution interpolated to the new resolution (see - :func:`interpolate_solution`). This drives the classic ``ns_array`` multi-grid and, - by also increasing ``mpol`` / ``ntor`` along the schedule, Fourier continuation, - entirely from Python. + :func:`interpolate_solution`); if ``restart_from`` is given, it seeds the first + step instead of a cold start. Args: - input: the target configuration. Its boundary is the final-resolution boundary; - each step truncates it. Schedule arrays default to the corresponding fields - of ``input``; ``mpol_array`` / ``ntor_array`` default to constant - ``input.mpol`` / ``input.ntor`` (i.e. the classic fixed-Fourier multi-grid). - ns_array, mpol_array, ntor_array, ftol_array, niter_array: per-step schedules. - All provided arrays must share one length (a length-1 array is broadcast). - **run_kwargs: forwarded to :func:`vmecpp.run` for every step (e.g. ``verbose``, - ``max_threads``). + input: the target configuration. Its boundary is the final-resolution + boundary; each step truncates or zero-pads it to that step's resolution. + magnetic_field, max_threads, verbose, restart_from: forwarded to + :func:`vmecpp.run` for every step (``restart_from`` only seeds the first). Returns: - The converged :class:`VmecOutput` at the final resolution. + The converged :class:`VmecOutput` at the final resolution, with ``input`` set + to the original (full-schedule) ``input`` argument. """ import vmecpp # noqa: PLC0415 (lazy import avoids a circular import) - ns_schedule = [int(x) for x in (input.ns_array if ns_array is None else ns_array)] + ns_schedule = [int(x) for x in input.ns_array] n_steps = len(ns_schedule) - if n_steps == 0: - msg = "ns_array must have at least one entry" - raise ValueError(msg) - - def _resolve(values: typing.Sequence[float] | None, default: list) -> list: - resolved = list(default) if values is None else list(values) - if len(resolved) == 1: - resolved = resolved * n_steps + + def _resolve(value: int | np.ndarray, name: str) -> list[int]: + if isinstance(value, int): + return [value] * n_steps + resolved = [int(x) for x in value] if len(resolved) != n_steps: msg = ( - f"continuation schedule length {len(resolved)} does not match " - f"ns_array length {n_steps}" + f"'{name}' has {len(resolved)} entries, but 'ns_array' has " + f"{n_steps}; a Fourier-resolution continuation schedule must have " + "one entry per ns_array step (or be a scalar, broadcast to every " + "step)." ) raise ValueError(msg) return resolved - mpol_schedule = _resolve(mpol_array, [int(input.mpol)] * n_steps) - ntor_schedule = _resolve(ntor_array, [int(input.ntor)] * n_steps) - ftol_schedule = _resolve(ftol_array, list(np.asarray(input.ftol_array))) - niter_schedule = _resolve(niter_array, list(np.asarray(input.niter_array))) + mpol_schedule = _resolve(input.mpol, "mpol") + ntor_schedule = _resolve(input.ntor, "ntor") + ftol_schedule = [float(x) for x in input.ftol_array] + niter_schedule = [int(x) for x in input.niter_array] - output: VmecOutput | None = None + output = restart_from for i in range(n_steps): step_input = _step_input( input, @@ -351,10 +354,14 @@ def _resolve(values: typing.Sequence[float] | None, default: list) -> list: ftol_schedule[i], niter_schedule[i], ) - if output is None: - output = vmecpp.run(step_input, **run_kwargs) - else: - guess = interpolate_solution(output, step_input) - output = vmecpp.run(step_input, restart_from=guess, **run_kwargs) + guess = None if output is None else interpolate_solution(output, step_input) + output = vmecpp.run( + step_input, + magnetic_field, + max_threads=max_threads, + verbose=verbose, + restart_from=guess, + ) + assert output is not None # n_steps >= 1, so the loop always assigns output - return output + return output.model_copy(update={"input": input}) diff --git a/src/vmecpp/simsopt_compat.py b/src/vmecpp/simsopt_compat.py index cddc26ab4..9778c8313 100644 --- a/src/vmecpp/simsopt_compat.py +++ b/src/vmecpp/simsopt_compat.py @@ -146,8 +146,9 @@ def __init__( # object, but the mpol/ntor values of either the vmec object # or the boundary surface object can be changed independently # by the user. + mpol, ntor = self._last_mpol_ntor(self.indata) mpol_for_surfacerzfourier, ntor_for_surfacerzfourier = ( - self._surface_rzfourier_resolution(self.indata.mpol, self.indata.ntor) + self._surface_rzfourier_resolution(mpol, ntor) ) self._boundary = SurfaceRZFourier.from_nphi_ntheta( nfp=self.indata.nfp, @@ -162,8 +163,8 @@ def __init__( # Transfer boundary shape data from indata to _boundary: vi = self.indata - for m in range(vi.mpol): - for n in range(2 * vi.ntor + 1): + for m in range(mpol): + for n in range(2 * ntor + 1): self._boundary.rc[m, n] = vi.rbc[m, n] self._boundary.zs[m, n] = vi.zbs[m, n] if vi.lasym: @@ -435,6 +436,18 @@ def boundary(self, boundary: SurfaceRZFourier) -> None: self.append_parent(boundary) self.need_to_run_code = True + @staticmethod + def _last_mpol_ntor(indata: vmecpp.VmecInput) -> tuple[int, int]: + """(indata.mpol, indata.ntor) as plain ints. + + VmecInput.mpol/.ntor may also be a Fourier-resolution continuation schedule (a + sequence); SIMSOPT's SurfaceRZFourier only has a single resolution, so the last + (finest, target) entry of the schedule is used. + """ + mpol = indata.mpol if isinstance(indata.mpol, int) else int(indata.mpol[-1]) + ntor = indata.ntor if isinstance(indata.ntor, int) else int(indata.ntor[-1]) + return mpol, ntor + @staticmethod def _surface_rzfourier_resolution(mpol: int, ntor: int) -> tuple[int, int]: # SurfaceRZFourier uses m up to mpol inclusive, unlike VMEC++. @@ -462,9 +475,8 @@ def set_indata(self) -> None: raise RuntimeError(msg) assert self.indata is not None vi = self.indata # Shorthand - target_mpol, target_ntor = self._surface_rzfourier_resolution( - self.indata.mpol, self.indata.ntor - ) + mpol, ntor = self._last_mpol_ntor(self.indata) + target_mpol, target_ntor = self._surface_rzfourier_resolution(mpol, ntor) boundary_RZFourier = self._resize_surface_rzfourier( self.boundary.to_RZFourier().copy(), target_mpol, @@ -483,8 +495,7 @@ def set_indata(self) -> None: zbc.fill(0.0) # Transfer boundary shape data from the surface object to VMEC: - ntor = self.indata.ntor - for m in range(self.indata.mpol): + for m in range(mpol): for n in range(2 * ntor + 1): vi.rbc[m, n] = boundary_RZFourier.get_rc(m, n - ntor) vi.zbs[m, n] = boundary_RZFourier.get_zs(m, n - ntor) diff --git a/tests/test_continuation.py b/tests/test_continuation.py index 41dbf20a0..2dbc4a3c3 100644 --- a/tests/test_continuation.py +++ b/tests/test_continuation.py @@ -125,42 +125,55 @@ def test_interpolate_fourier_pad_and_truncate(solovev_output: vmecpp.VmecOutput) # --- continuation driver ----------------------------------------------------- -def test_single_step_continuation_matches_direct_run(cma_input: vmecpp.VmecInput): - """A single-resolution continuation invokes exactly one solve and must be bit-for- - bit identical to calling ``run`` on the same single-grid input.""" - ns_final = int(np.asarray(cma_input.ns_array)[-1]) - - single_grid = cma_input.model_copy(deep=True) - single_grid.ns_array = np.asarray([ns_final], dtype=np.int64) - single_grid.ftol_array = np.asarray([1e-12], dtype=float) - single_grid.niter_array = np.asarray([60000], dtype=np.int64) - direct = vmecpp.run(single_grid, verbose=False, max_threads=1) - - continued = vmecpp.run_continuation( - cma_input, - ns_array=[ns_final], - ftol_array=[1e-12], - niter_array=[60000], - verbose=False, - max_threads=1, +def test_mpol_ntor_length_one_sequence_collapses_to_scalar( + cma_input: vmecpp.VmecInput, +): + """A length-1 mpol/ntor sequence is equivalent to a scalar and is stored as a plain + int, so it takes the direct (non-continuation) code path. + + This coercion runs during validation (construction / model_validate), like all of + VmecInput's other validators; it goes through vmecpp.VmecInput.model_validate here + rather than a bare attribute assignment for that reason. + """ + data = cma_input.model_dump(mode="json") + data["mpol"] = [int(cma_input.mpol)] + data["ntor"] = [int(cma_input.ntor)] + single_step = vmecpp.VmecInput.model_validate(data) + assert isinstance(single_step.mpol, int) + assert isinstance(single_step.ntor, int) + assert single_step.mpol == cma_input.mpol + assert single_step.ntor == cma_input.ntor + + +def test_mpol_length_mismatch_raises(cma_input: vmecpp.VmecInput): + """A mpol/ntor schedule that doesn't match ns_array's length is rejected with a + clear error, rather than silently misaligning steps.""" + mismatched = cma_input.model_copy(deep=True) + mismatched.ns_array = np.asarray([15, 31], dtype=np.int64) + mismatched.ftol_array = np.asarray([1e-8, 1e-10], dtype=float) + mismatched.niter_array = np.asarray([100, 100], dtype=np.int64) + mismatched.mpol = np.asarray([3, 4, 5], dtype=np.int64) # 3 entries, ns_array has 2 + + with pytest.raises(ValueError, match="mpol"): + vmecpp.run(mismatched, verbose=False, max_threads=1) + + +def test_ns_only_continuation_reproduces_direct_multigrid( + cma_input: vmecpp.VmecInput, cma_direct: vmecpp.VmecOutput +): + """A continuation schedule with a constant (array-valued) mpol/ntor reaches the same + equilibrium as the C++ multi-grid: identical resolution, a converged force balance, + a bit-identical plasma volume, and geometry agreeing at the convergence level.""" + n_steps = len(np.asarray(cma_input.ns_array)) + continuation_input = cma_input.model_copy(deep=True) + continuation_input.mpol = np.asarray( + [int(cma_input.mpol)] * n_steps, dtype=np.int64 ) - for field in ("rmnc", "zmns", "lmns_full"): - np.testing.assert_array_equal( - np.asarray(getattr(continued.wout, field)), - np.asarray(getattr(direct.wout, field)), - ) - np.testing.assert_array_equal( - np.asarray(continued.wout.iotaf), np.asarray(direct.wout.iotaf) + continuation_input.ntor = np.asarray( + [int(cma_input.ntor)] * n_steps, dtype=np.int64 ) - -def test_run_continuation_reproduces_direct( - cma_input: vmecpp.VmecInput, cma_direct: vmecpp.VmecOutput -): - """The default Python continuation reaches the same equilibrium as the C++ - multi-grid: identical resolution, a converged force balance, a bit-identical - plasma volume, and geometry agreeing at the convergence level.""" - continued = vmecpp.run_continuation(cma_input, verbose=False, max_threads=1) + continued = vmecpp.run(continuation_input, verbose=False, max_threads=1) assert int(continued.wout.ns) == int(cma_direct.wout.ns) # Converged to a force residual comparable to the direct solve. @@ -184,14 +197,23 @@ def test_continuation_agreement_tightens_with_ftol(cma_input: vmecpp.VmecInput): """As the force-balance tolerance tightens, the continuation and the direct multi- grid converge toward the same geometry -- the signature of a shared equilibrium rather than a defect in the continuation.""" + n_steps = len(np.asarray(cma_input.ns_array)) def max_geometry_diff(ftol_array: list[float]) -> float: reference = cma_input.model_copy(deep=True) reference.ftol_array = np.asarray(ftol_array, dtype=float) direct = vmecpp.run(reference, verbose=False, max_threads=1) - continued = vmecpp.run_continuation( - cma_input, ftol_array=ftol_array, verbose=False, max_threads=1 + + continuation_input = cma_input.model_copy(deep=True) + continuation_input.mpol = np.asarray( + [int(cma_input.mpol)] * n_steps, dtype=np.int64 ) + continuation_input.ntor = np.asarray( + [int(cma_input.ntor)] * n_steps, dtype=np.int64 + ) + continuation_input.ftol_array = np.asarray(ftol_array, dtype=float) + continued = vmecpp.run(continuation_input, verbose=False, max_threads=1) + return float( np.max( np.abs(np.asarray(direct.wout.rmnc) - np.asarray(continued.wout.rmnc)) @@ -214,16 +236,18 @@ def test_fourier_continuation_converges( ftol_final = float(np.asarray(cma_input.ftol_array)[-1]) niter_final = int(np.asarray(cma_input.niter_array)[-1]) - continued = vmecpp.run_continuation( - cma_input, - ns_array=[ns_final, ns_final], - mpol_array=[max(2, mpol_final - 2), mpol_final], - ntor_array=[ntor, ntor], - ftol_array=[ftol_final, ftol_final], - niter_array=[niter_final, niter_final], - verbose=False, - max_threads=1, + continuation_input = cma_input.model_copy(deep=True) + continuation_input.ns_array = np.asarray([ns_final, ns_final], dtype=np.int64) + continuation_input.mpol = np.asarray( + [max(2, mpol_final - 2), mpol_final], dtype=np.int64 ) + continuation_input.ntor = ntor # scalar broadcasts to every step + continuation_input.ftol_array = np.asarray([ftol_final, ftol_final], dtype=float) + continuation_input.niter_array = np.asarray( + [niter_final, niter_final], dtype=np.int64 + ) + + continued = vmecpp.run(continuation_input, verbose=False, max_threads=1) assert int(continued.wout.mpol) == mpol_final assert _final_force_residual(continued) < 1e-5 @@ -236,3 +260,36 @@ def test_fourier_continuation_converges( atol=4e-3, rtol=1e-2, ) + + +def test_fourier_continuation_hot_restarts_first_step_from_restart_from( + cma_input: vmecpp.VmecInput, +): + """restart_from seeds the first continuation step (interpolated to its resolution) + instead of a cold start, when input.mpol/.ntor is a sequence.""" + ns_final = int(np.asarray(cma_input.ns_array)[-1]) + mpol_final = int(cma_input.mpol) + ntor = int(cma_input.ntor) + + warm_start_input = cma_input.model_copy(deep=True) + warm_start_input.ns_array = np.asarray([ns_final], dtype=np.int64) + warm_start = vmecpp.run(warm_start_input, verbose=False, max_threads=1) + + continuation_input = cma_input.model_copy(deep=True) + continuation_input.ns_array = np.asarray([ns_final, ns_final], dtype=np.int64) + continuation_input.mpol = np.asarray( + [max(2, mpol_final - 2), mpol_final], dtype=np.int64 + ) + continuation_input.ntor = ntor + continuation_input.ftol_array = np.asarray([1e-8, 1e-10], dtype=float) + continuation_input.niter_array = np.asarray([2000, 2000], dtype=np.int64) + + continued = vmecpp.run( + continuation_input, + verbose=False, + max_threads=1, + restart_from=warm_start, + ) + + assert int(continued.wout.mpol) == mpol_final + assert _final_force_residual(continued) < 1e-5 From 271f3bf3e621c39c26235980a5df89062cf120fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Jura=C5=A1i=C4=87?= <166746189+jurasic-pf@users.noreply.github.com> Date: Thu, 9 Jul 2026 20:55:55 +0200 Subject: [PATCH 05/12] Fix transform hot-loop regression from Eigen3 migration (#621) The Eigen3 migration (#410) and hot-loop rework (#454) replaced the fused scalar poloidal accumulation in the toroidal transforms with per-quantity Eigen .dot() calls, and the FFT path additionally allocates two Eigen vectors per innermost (m,k) iteration via .eval(). On the short theta axis (nThetaReduced ~9-16) and small ntor+1 this is a pessimization: benchmark-runs history shows ToroidalForcesToFourier regressed ~2x from the pre-#410 baseline and never recovered, including at the flagship 12x12 FFT size. Two changes: - dft_toroidal.cc: restore the pre-#410 fused-scalar-loop DFT code verbatim (single pass over theta reading each basis value once; the original "auto-vectorize was a pessimization" note is retained). Fixes the DFT-fallback resolutions and the fftx-disabled build. - fft_toroidal.cc: the FFT path only replaces the toroidal direction; its poloidal fill kept the .dot()+.eval() pattern. Fuse it into one allocation-free scalar pass. Measured (same-machine A/B, --config=opt, OMP=1): 12x12 FFT forces 1.57x faster (1.00e-3 -> 6.35e-4), 6x8 neutral. The c2r FourierToReal output scatter already uses plain segment += and is left unchanged. fft_toroidal_test and vmec_test pass. CI benchmarks will confirm the recovery and inform whether any FFT shapes should still fall back to DFT. Co-authored-by: Claude Opus 4.8 (1M context) --- .../vmec/ideal_mhd_model/dft_toroidal.cc | 388 +++++++++--------- .../vmec/ideal_mhd_model/fft_toroidal.cc | 137 ++++--- 2 files changed, 267 insertions(+), 258 deletions(-) diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc index be9acf83c..d2e92218e 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc @@ -51,95 +51,75 @@ void ForcesToFourier3DSymmFastPoloidal( const auto& fzcon = m_even ? d.fzcon_e : d.fzcon_o; for (int k = 0; k < s.nZeta; ++k) { + double rmkcc = 0.0; + double rmkcc_n = 0.0; + double rmkss = 0.0; + double rmkss_n = 0.0; + double zmksc = 0.0; + double zmksc_n = 0.0; + double zmkcs = 0.0; + double zmkcs_n = 0.0; + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; const int idx_ml_base = m * s.nThetaReduced; - // Vectorized poloidal loop using Eigen operations - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); - - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - auto crmn_seg = Eigen::Map( - crmn.data() + idx_kl_base, s.nThetaReduced); - auto czmn_seg = Eigen::Map( - czmn.data() + idx_kl_base, s.nThetaReduced); - auto armn_seg = Eigen::Map( - armn.data() + idx_kl_base, s.nThetaReduced); - auto azmn_seg = Eigen::Map( - azmn.data() + idx_kl_base, s.nThetaReduced); - auto brmn_seg = Eigen::Map( - brmn.data() + idx_kl_base, s.nThetaReduced); - auto bzmn_seg = Eigen::Map( - bzmn.data() + idx_kl_base, s.nThetaReduced); - auto frcon_seg = Eigen::Map( - frcon.data() + idx_kl_base, s.nThetaReduced); - auto fzcon_seg = Eigen::Map( - fzcon.data() + idx_kl_base, s.nThetaReduced); - - double lmksc = blmn_seg.dot(cosmumi_seg); - double lmkcs = blmn_seg.dot(sinmumi_seg); - double lmkcs_n = -clmn_seg.dot(cosmui_seg); - double lmksc_n = -clmn_seg.dot(sinmui_seg); - - double rmkcc_n = -crmn_seg.dot(cosmui_seg); - double zmkcs_n = -czmn_seg.dot(cosmui_seg); - - double rmkss_n = -crmn_seg.dot(sinmui_seg); - double zmksc_n = -czmn_seg.dot(sinmui_seg); - - // Assemble effective R and Z forces from MHD and spectral condensation - // contributions. Materialize to avoid re-evaluation in each dot - // product. - // Per-thread scratch reused across iterations instead of a heap - // temporary in this innermost loop; still materialized once and then - // used in the two dot products below. - thread_local Eigen::VectorXd tempR_seg; - thread_local Eigen::VectorXd tempZ_seg; - tempR_seg = armn_seg + xmpq[m] * frcon_seg; - tempZ_seg = azmn_seg + xmpq[m] * fzcon_seg; - - double rmkcc = tempR_seg.dot(cosmui_seg) + brmn_seg.dot(sinmumi_seg); - double rmkss = tempR_seg.dot(sinmui_seg) + brmn_seg.dot(cosmumi_seg); - double zmksc = tempZ_seg.dot(sinmui_seg) + bzmn_seg.dot(cosmumi_seg); - double zmkcs = tempZ_seg.dot(cosmui_seg) + bzmn_seg.dot(sinmumi_seg); - - // Vectorized toroidal scatter: segment ops replace scalar n-loop - const int ntorp1 = s.ntor + 1; - const int idx_mn_base = ((jF - rp.nsMinF) * s.mpol + m) * ntorp1; - const int idx_kn_base = k * (s.nnyq2 + 1); - - auto cosnv_seg = fb.cosnv.segment(idx_kn_base, ntorp1); - auto sinnv_seg = fb.sinnv.segment(idx_kn_base, ntorp1); - auto cosnvn_seg = fb.cosnvn.segment(idx_kn_base, ntorp1); - auto sinnvn_seg = fb.sinnvn.segment(idx_kn_base, ntorp1); - - Eigen::Map frcc_seg( - m_physical_forces.frcc.data() + idx_mn_base, ntorp1); - Eigen::Map frss_seg( - m_physical_forces.frss.data() + idx_mn_base, ntorp1); - Eigen::Map fzsc_seg( - m_physical_forces.fzsc.data() + idx_mn_base, ntorp1); - Eigen::Map fzcs_seg( - m_physical_forces.fzcs.data() + idx_mn_base, ntorp1); - - frcc_seg += rmkcc * cosnv_seg + rmkcc_n * sinnvn_seg; - frss_seg += rmkss * sinnv_seg + rmkss_n * cosnvn_seg; - fzsc_seg += zmksc * cosnv_seg + zmksc_n * sinnvn_seg; - fzcs_seg += zmkcs * sinnv_seg + zmkcs_n * cosnvn_seg; - - if (jMinL <= jF) { - Eigen::Map flsc_seg( - m_physical_forces.flsc.data() + idx_mn_base, ntorp1); - Eigen::Map flcs_seg( - m_physical_forces.flcs.data() + idx_mn_base, ntorp1); - flsc_seg += lmksc * cosnv_seg + lmksc_n * sinnvn_seg; - flcs_seg += lmkcs * sinnv_seg + lmkcs_n * cosnvn_seg; - } + // NOTE: nThetaReduced is usually pretty small, 9 for cma.json + // and 16 for w7x_ref_167_12_12.json, so in our benchmark forcing + // the compiler to auto-vectorize this loop was a pessimization. + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; // --> flsc (no A) + lmkcs += blmn[idx_kl] * sinmumi; // --> flcs + lmkcs_n -= clmn[idx_kl] * cosmui; // --> flcs + lmksc_n -= clmn[idx_kl] * sinmui; // --> flsc + + rmkcc_n -= crmn[idx_kl] * cosmui; // --> frcc + zmkcs_n -= czmn[idx_kl] * cosmui; // --> fzcs + + rmkss_n -= crmn[idx_kl] * sinmui; // --> frss + zmksc_n -= czmn[idx_kl] * sinmui; // --> fzsc + + // assemble effective R and Z forces from MHD and spectral + // condensation contributions + const double tempR = armn[idx_kl] + xmpq[m] * frcon[idx_kl]; + const double tempZ = azmn[idx_kl] + xmpq[m] * fzcon[idx_kl]; + + rmkcc += tempR * cosmui + brmn[idx_kl] * sinmumi; // --> frcc + rmkss += tempR * sinmui + brmn[idx_kl] * cosmumi; // --> frss + zmksc += tempZ * sinmui + bzmn[idx_kl] * cosmumi; // --> fzsc + zmkcs += tempZ * cosmui + bzmn[idx_kl] * sinmumi; // --> fzcs + } // l + + for (int n = 0; n < s.ntor + 1; ++n) { + const int idx_mn = ((jF - rp.nsMinF) * s.mpol + m) * (s.ntor + 1) + n; + const int idx_kn = k * (s.nnyq2 + 1) + n; + + const double cosnv = fb.cosnv[idx_kn]; + const double sinnv = fb.sinnv[idx_kn]; + const double cosnvn = fb.cosnvn[idx_kn]; + const double sinnvn = fb.sinnvn[idx_kn]; + + m_physical_forces.frcc[idx_mn] += rmkcc * cosnv + rmkcc_n * sinnvn; + m_physical_forces.frss[idx_mn] += rmkss * sinnv + rmkss_n * cosnvn; + m_physical_forces.fzsc[idx_mn] += zmksc * cosnv + zmksc_n * sinnvn; + m_physical_forces.fzcs[idx_mn] += zmkcs * sinnv + zmkcs_n * cosnvn; + + if (jMinL <= jF) { + m_physical_forces.flsc[idx_mn] += lmksc * cosnv + lmksc_n * sinnvn; + m_physical_forces.flcs[idx_mn] += lmkcs * sinnv + lmkcs_n * cosnvn; + } + } // n } // k } // m } // jF @@ -154,42 +134,41 @@ void ForcesToFourier3DSymmFastPoloidal( const auto& clmn = m_even ? d.clmn_e : d.clmn_o; for (int k = 0; k < s.nZeta; ++k) { + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; const int idx_ml_base = m * s.nThetaReduced; - // Vectorized poloidal loop using Eigen operations - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); - - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - - double lmksc = blmn_seg.dot(cosmumi_seg); - double lmkcs = blmn_seg.dot(sinmumi_seg); - double lmkcs_n = -clmn_seg.dot(cosmui_seg); - double lmksc_n = -clmn_seg.dot(sinmui_seg); - - // Vectorized toroidal scatter for lambda-only section - const int ntorp1 = s.ntor + 1; - const int idx_mn_base = ((jF - rp.nsMinF) * s.mpol + m) * ntorp1; - const int idx_kn_base = k * (s.nnyq2 + 1); - - auto cosnv_seg = fb.cosnv.segment(idx_kn_base, ntorp1); - auto sinnv_seg = fb.sinnv.segment(idx_kn_base, ntorp1); - auto cosnvn_seg = fb.cosnvn.segment(idx_kn_base, ntorp1); - auto sinnvn_seg = fb.sinnvn.segment(idx_kn_base, ntorp1); - - Eigen::Map flsc_seg( - m_physical_forces.flsc.data() + idx_mn_base, ntorp1); - Eigen::Map flcs_seg( - m_physical_forces.flcs.data() + idx_mn_base, ntorp1); - - flsc_seg += lmksc * cosnv_seg + lmksc_n * sinnvn_seg; - flcs_seg += lmkcs * sinnv_seg + lmkcs_n * cosnvn_seg; + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; // --> flsc (no A) + lmkcs += blmn[idx_kl] * sinmumi; // --> flcs + lmkcs_n -= clmn[idx_kl] * cosmui; // --> flcs + lmksc_n -= clmn[idx_kl] * sinmui; // --> flsc + } // l + + for (int n = 0; n < s.ntor + 1; ++n) { + const int idx_mn = ((jF - rp.nsMinF) * s.mpol + m) * (s.ntor + 1) + n; + const int idx_kn = k * (s.nnyq2 + 1) + n; + + const double cosnv = fb.cosnv[idx_kn]; + const double sinnv = fb.sinnv[idx_kn]; + const double cosnvn = fb.cosnvn[idx_kn]; + const double sinnvn = fb.sinnvn[idx_kn]; + + m_physical_forces.flsc[idx_mn] += lmksc * cosnv + lmksc_n * sinnvn; + m_physical_forces.flcs[idx_mn] += lmkcs * sinnv + lmkcs_n * cosnvn; + } // n } // k } // m } // jF @@ -257,97 +236,100 @@ void FourierToReal3DSymmFastPoloidal(const FourierGeometry& physical_x, } for (int k = 0; k < s.nZeta; ++k) { - // INVERSE TRANSFORM IN N-ZETA, FOR FIXED M - // Vectorized toroidal accumulation loop - const int idx_kn_base = k * (s.nnyq2 + 1); - const int idx_mn_base = ((jF - nsMinF1) * s.mpol + m) * (s.ntor + 1); - - auto cosnv_seg = fb.cosnv.segment(idx_kn_base, s.ntor + 1); - auto sinnv_seg = fb.sinnv.segment(idx_kn_base, s.ntor + 1); - auto sinnvn_seg = fb.sinnvn.segment(idx_kn_base, s.ntor + 1); - auto cosnvn_seg = fb.cosnvn.segment(idx_kn_base, s.ntor + 1); - - auto rmncc_seg = Eigen::Map( - physical_x.rmncc.data() + idx_mn_base, s.ntor + 1); - auto rmnss_seg = Eigen::Map( - physical_x.rmnss.data() + idx_mn_base, s.ntor + 1); - auto zmnsc_seg = Eigen::Map( - physical_x.zmnsc.data() + idx_mn_base, s.ntor + 1); - auto zmncs_seg = Eigen::Map( - physical_x.zmncs.data() + idx_mn_base, s.ntor + 1); - auto lmnsc_seg = Eigen::Map( - physical_x.lmnsc.data() + idx_mn_base, s.ntor + 1); - auto lmncs_seg = Eigen::Map( - physical_x.lmncs.data() + idx_mn_base, s.ntor + 1); - - double rmkcc = rmncc_seg.dot(cosnv_seg); - double rmkcc_n = rmncc_seg.dot(sinnvn_seg); - double rmkss = rmnss_seg.dot(sinnv_seg); - double rmkss_n = rmnss_seg.dot(cosnvn_seg); - double zmksc = zmnsc_seg.dot(cosnv_seg); - double zmksc_n = zmnsc_seg.dot(sinnvn_seg); - double zmkcs = zmncs_seg.dot(sinnv_seg); - double zmkcs_n = zmncs_seg.dot(cosnvn_seg); - double lmksc = lmnsc_seg.dot(cosnv_seg); - double lmksc_n = lmnsc_seg.dot(sinnvn_seg); - double lmkcs = lmncs_seg.dot(sinnv_seg); - double lmkcs_n = lmncs_seg.dot(cosnvn_seg); + double rmkcc = 0.0; + double rmkcc_n = 0.0; + double rmkss = 0.0; + double rmkss_n = 0.0; + double zmksc = 0.0; + double zmksc_n = 0.0; + double zmkcs = 0.0; + double zmkcs_n = 0.0; + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + + for (int n = 0; n < s.ntor + 1; ++n) { + // INVERSE TRANSFORM IN N-ZETA, FOR FIXED M + + const int idx_kn = k * (s.nnyq2 + 1) + n; + + double cosnv = fb.cosnv[idx_kn]; + double sinnv = fb.sinnv[idx_kn]; + double sinnvn = fb.sinnvn[idx_kn]; + double cosnvn = fb.cosnvn[idx_kn]; + + int idx_mn = ((jF - nsMinF1) * s.mpol + m) * (s.ntor + 1) + n; + + rmkcc += physical_x.rmncc[idx_mn] * cosnv; + rmkcc_n += physical_x.rmncc[idx_mn] * sinnvn; + rmkss += physical_x.rmnss[idx_mn] * sinnv; + rmkss_n += physical_x.rmnss[idx_mn] * cosnvn; + zmksc += physical_x.zmnsc[idx_mn] * cosnv; + zmksc_n += physical_x.zmnsc[idx_mn] * sinnvn; + zmkcs += physical_x.zmncs[idx_mn] * sinnv; + zmkcs_n += physical_x.zmncs[idx_mn] * cosnvn; + lmksc += physical_x.lmnsc[idx_mn] * cosnv; + lmksc_n += physical_x.lmnsc[idx_mn] * sinnvn; + lmkcs += physical_x.lmncs[idx_mn] * sinnv; + lmkcs_n += physical_x.lmncs[idx_mn] * cosnvn; + } // n // INVERSE TRANSFORM IN M-THETA, FOR ALL RADIAL, ZETA VALUES const int idx_kl_base = ((jF - nsMinF1) * s.nZeta + k) * s.nThetaEff; - // Vectorized poloidal loops using Eigen operations - auto sinmum_seg = fb.sinmum.segment(idx_ml_base, s.nThetaReduced); - auto cosmum_seg = fb.cosmum.segment(idx_ml_base, s.nThetaReduced); - - auto ru_seg = Eigen::Map(ru.data() + idx_kl_base, - s.nThetaReduced); - auto zu_seg = Eigen::Map(zu.data() + idx_kl_base, - s.nThetaReduced); - auto lu_seg = Eigen::Map(lu.data() + idx_kl_base, - s.nThetaReduced); - - // NOTE: element-wise multiplication - ru_seg += rmkcc * sinmum_seg + rmkss * cosmum_seg; - zu_seg += zmksc * cosmum_seg + zmkcs * sinmum_seg; - lu_seg += lmksc * cosmum_seg + lmkcs * sinmum_seg; - - auto cosmu_seg = fb.cosmu.segment(idx_ml_base, s.nThetaReduced); - auto sinmu_seg = fb.sinmu.segment(idx_ml_base, s.nThetaReduced); - - auto rv_seg = Eigen::Map(rv.data() + idx_kl_base, - s.nThetaReduced); - auto zv_seg = Eigen::Map(zv.data() + idx_kl_base, - s.nThetaReduced); - auto lv_seg = Eigen::Map(lv.data() + idx_kl_base, - s.nThetaReduced); - - // NOTE: element-wise multiplication - rv_seg += rmkcc_n * cosmu_seg + rmkss_n * sinmu_seg; - zv_seg += zmksc_n * sinmu_seg + zmkcs_n * cosmu_seg; - // it is here that lv gets a negative sign! - lv_seg -= lmksc_n * sinmu_seg + lmkcs_n * cosmu_seg; - - auto r1_seg = Eigen::Map(r1.data() + idx_kl_base, - s.nThetaReduced); - auto z1_seg = Eigen::Map(z1.data() + idx_kl_base, - s.nThetaReduced); - - r1_seg += rmkcc * cosmu_seg + rmkss * sinmu_seg; - z1_seg += zmksc * sinmu_seg + zmkcs * cosmu_seg; + // the loop over l is split to help compiler auto-vectorization + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_ml = idx_ml_base + l; - if (nsMinF <= jF && jF < r.nsMaxFIncludingLcfs) { - // spectral condensation is local per flux surface - const int idx_con_base = ((jF - nsMinF) * s.nZeta + k) * s.nThetaEff; + const double sinmum = fb.sinmum[idx_ml]; + const double cosmum = fb.cosmum[idx_ml]; + + const int idx_kl = idx_kl_base + l; + ru[idx_kl] += rmkcc * sinmum + rmkss * cosmum; + zu[idx_kl] += zmksc * cosmum + zmkcs * sinmum; + lu[idx_kl] += lmksc * cosmum + lmkcs * sinmum; + } // l + + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; - auto rCon_seg = Eigen::Map( - m_geometry.rCon.data() + idx_con_base, s.nThetaReduced); - auto zCon_seg = Eigen::Map( - m_geometry.zCon.data() + idx_con_base, s.nThetaReduced); + const double cosmu = fb.cosmu[idx_ml]; + const double sinmu = fb.sinmu[idx_ml]; + rv[idx_kl] += rmkcc_n * cosmu + rmkss_n * sinmu; + zv[idx_kl] += zmksc_n * sinmu + zmkcs_n * cosmu; + // it is here that lv gets a negative sign! + lv[idx_kl] -= lmksc_n * sinmu + lmkcs_n * cosmu; + } // l - rCon_seg += (rmkcc * cosmu_seg + rmkss * sinmu_seg) * con_factor; - zCon_seg += (zmksc * sinmu_seg + zmkcs * cosmu_seg) * con_factor; - } + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_ml = idx_ml_base + l; + + const double cosmu = fb.cosmu[idx_ml]; + const double sinmu = fb.sinmu[idx_ml]; + + const int idx_kl = idx_kl_base + l; + + r1[idx_kl] += rmkcc * cosmu + rmkss * sinmu; + z1[idx_kl] += zmksc * sinmu + zmkcs * cosmu; + } // l + + if (nsMinF <= jF && jF < r.nsMaxFIncludingLcfs) { + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_ml = idx_ml_base + l; + const double cosmu = fb.cosmu[idx_ml]; + const double sinmu = fb.sinmu[idx_ml]; + + // spectral condensation is local per flux surface + // --> no need for numFull1 + const int idx_con = ((jF - nsMinF) * s.nZeta + k) * s.nThetaEff + l; + m_geometry.rCon[idx_con] += + (rmkcc * cosmu + rmkss * sinmu) * con_factor; + m_geometry.zCon[idx_con] += + (zmksc * sinmu + zmkcs * cosmu) * con_factor; + } + } // l } // k } // m } // j diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc index 8b845c752..4323a1c30 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc @@ -459,10 +459,6 @@ void ForcesToFourier3DSymmFastPoloidalFft( const auto& fzcon = m_even ? d.fzcon_e : d.fzcon_o; const int idx_ml_base = m * s.nThetaReduced; - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); double* rmkcc_buf = in_slot(m, kRmkcc); double* rmkss_buf = in_slot(m, kRmkss); @@ -477,47 +473,65 @@ void ForcesToFourier3DSymmFastPoloidalFft( double* lmksc_n_buf = in_slot(m, kLmkscN); double* lmkcs_n_buf = in_slot(m, kLmkcsN); + const double xmpq_m = xmpq[m]; for (int k = 0; k < s.nZeta; ++k) { const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - auto crmn_seg = Eigen::Map( - crmn.data() + idx_kl_base, s.nThetaReduced); - auto czmn_seg = Eigen::Map( - czmn.data() + idx_kl_base, s.nThetaReduced); - auto armn_seg = Eigen::Map( - armn.data() + idx_kl_base, s.nThetaReduced); - auto azmn_seg = Eigen::Map( - azmn.data() + idx_kl_base, s.nThetaReduced); - auto brmn_seg = Eigen::Map( - brmn.data() + idx_kl_base, s.nThetaReduced); - auto bzmn_seg = Eigen::Map( - bzmn.data() + idx_kl_base, s.nThetaReduced); - auto frcon_seg = Eigen::Map( - frcon.data() + idx_kl_base, s.nThetaReduced); - auto fzcon_seg = Eigen::Map( - fzcon.data() + idx_kl_base, s.nThetaReduced); - - lmksc_buf[k] = blmn_seg.dot(cosmumi_seg); - lmkcs_buf[k] = blmn_seg.dot(sinmumi_seg); - lmkcs_n_buf[k] = -clmn_seg.dot(cosmui_seg); - lmksc_n_buf[k] = -clmn_seg.dot(sinmui_seg); - - rmkcc_n_buf[k] = -crmn_seg.dot(cosmui_seg); - zmkcs_n_buf[k] = -czmn_seg.dot(cosmui_seg); - rmkss_n_buf[k] = -crmn_seg.dot(sinmui_seg); - zmksc_n_buf[k] = -czmn_seg.dot(sinmui_seg); - - const Eigen::VectorXd tempR = (armn_seg + xmpq[m] * frcon_seg).eval(); - const Eigen::VectorXd tempZ = (azmn_seg + xmpq[m] * fzcon_seg).eval(); - - rmkcc_buf[k] = tempR.dot(cosmui_seg) + brmn_seg.dot(sinmumi_seg); - rmkss_buf[k] = tempR.dot(sinmui_seg) + brmn_seg.dot(cosmumi_seg); - zmksc_buf[k] = tempZ.dot(sinmui_seg) + bzmn_seg.dot(cosmumi_seg); - zmkcs_buf[k] = tempZ.dot(cosmui_seg) + bzmn_seg.dot(sinmumi_seg); + double rmkcc = 0.0; + double rmkcc_n = 0.0; + double rmkss = 0.0; + double rmkss_n = 0.0; + double zmksc = 0.0; + double zmksc_n = 0.0; + double zmkcs = 0.0; + double zmkcs_n = 0.0; + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + + // Fused poloidal reduction, matching ForcesToFourier3DSymmFastPoloidal: + // one pass over the short theta axis, no per-iteration heap temporary. + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; + lmkcs += blmn[idx_kl] * sinmumi; + lmkcs_n -= clmn[idx_kl] * cosmui; + lmksc_n -= clmn[idx_kl] * sinmui; + + rmkcc_n -= crmn[idx_kl] * cosmui; + zmkcs_n -= czmn[idx_kl] * cosmui; + rmkss_n -= crmn[idx_kl] * sinmui; + zmksc_n -= czmn[idx_kl] * sinmui; + + const double tempR = armn[idx_kl] + xmpq_m * frcon[idx_kl]; + const double tempZ = azmn[idx_kl] + xmpq_m * fzcon[idx_kl]; + + rmkcc += tempR * cosmui + brmn[idx_kl] * sinmumi; + rmkss += tempR * sinmui + brmn[idx_kl] * cosmumi; + zmksc += tempZ * sinmui + bzmn[idx_kl] * cosmumi; + zmkcs += tempZ * cosmui + bzmn[idx_kl] * sinmumi; + } // l + + rmkcc_buf[k] = rmkcc; + rmkss_buf[k] = rmkss; + rmkcc_n_buf[k] = rmkcc_n; + rmkss_n_buf[k] = rmkss_n; + zmksc_buf[k] = zmksc; + zmkcs_buf[k] = zmkcs; + zmksc_n_buf[k] = zmksc_n; + zmkcs_n_buf[k] = zmkcs_n; + lmksc_buf[k] = lmksc; + lmkcs_buf[k] = lmkcs; + lmksc_n_buf[k] = lmksc_n; + lmkcs_n_buf[k] = lmkcs_n; } // k } // m (fill) @@ -590,10 +604,6 @@ void ForcesToFourier3DSymmFastPoloidalFft( const auto& clmn = m_even ? d.clmn_e : d.clmn_o; const int idx_ml_base = m * s.nThetaReduced; - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); double* lmksc_buf = in_slot(m, kLmksc); double* lmkcs_buf = in_slot(m, kLmkcs); @@ -602,15 +612,32 @@ void ForcesToFourier3DSymmFastPoloidalFft( for (int k = 0; k < s.nZeta; ++k) { const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - - lmksc_buf[k] = blmn_seg.dot(cosmumi_seg); - lmkcs_buf[k] = blmn_seg.dot(sinmumi_seg); - lmkcs_n_buf[k] = -clmn_seg.dot(cosmui_seg); - lmksc_n_buf[k] = -clmn_seg.dot(sinmui_seg); + + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + + // Fused poloidal reduction (see main-loop note). + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; + lmkcs += blmn[idx_kl] * sinmumi; + lmkcs_n -= clmn[idx_kl] * cosmui; + lmksc_n -= clmn[idx_kl] * sinmui; + } // l + + lmksc_buf[k] = lmksc; + lmkcs_buf[k] = lmkcs; + lmksc_n_buf[k] = lmksc_n; + lmkcs_n_buf[k] = lmkcs_n; } // k } // m (fill) From ec3ee0919e386221e88b01d59bfbd0a542df814b Mon Sep 17 00:00:00 2001 From: CharlesCNorton Date: Thu, 9 Jul 2026 15:17:55 -0400 Subject: [PATCH 06/12] Fix typos in documentation (#627) Fix typos in docs --- examples/data/README.md | 4 ++-- src/vmecpp/cpp/docker/tsan/README.md | 2 +- src/vmecpp/cpp/docker/tsan/background_info.md | 16 ++++++++-------- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/examples/data/README.md b/examples/data/README.md index 00efb23fd..543b37f78 100644 --- a/examples/data/README.md +++ b/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 diff --git a/src/vmecpp/cpp/docker/tsan/README.md b/src/vmecpp/cpp/docker/tsan/README.md index 33bd6bbe2..c60eea9a8 100644 --- a/src/vmecpp/cpp/docker/tsan/README.md +++ b/src/vmecpp/cpp/docker/tsan/README.md @@ -1,6 +1,6 @@ # OpenMP-aware ThreadSanitizer -This is a tutorial how to build `clang` and `libomp` with ThreadSanitizer (TSan) +This is a tutorial on how to build `clang` and `libomp` with ThreadSanitizer (TSan) support for being able to check OpenMP-parallelized C/C++ code. ## Perform ThreadSanitizer runs on code in the Proxima Fusion monorepo diff --git a/src/vmecpp/cpp/docker/tsan/background_info.md b/src/vmecpp/cpp/docker/tsan/background_info.md index 60035b0dc..0a3b9a58d 100644 --- a/src/vmecpp/cpp/docker/tsan/background_info.md +++ b/src/vmecpp/cpp/docker/tsan/background_info.md @@ -28,7 +28,7 @@ found here: - https://www.vi-hps.org/cms/upload/material/tw30/Archer.pdf (slide 8 and onwards) Turns out, this is **deprecated** as well and Archer was integrated into the -LLVM project in the mean time. +LLVM project in the meantime. This seems to be the **state of things as of Oct 2023**. Indeed, Archer can be found in the `llvm` source tree: @@ -46,7 +46,7 @@ https://packages.ubuntu.com/jammy/clang ## Setting up a first test case TL;DR: This is an example commonly presented as a case for data races. As will -be seen when running this example, the errornous output due to a data race is +be seen when running this example, the erroneous output due to a data race is not reliably reproduced. Thus, skip this section if you are looking for an example that breaks reliably. @@ -223,7 +223,7 @@ int main(int argc, char** argv) { } ``` -Compile it (for now without TSan to not clobber the console when things to +Compile it (for now without TSan to not clobber the console when things go sideways): ```bash @@ -313,7 +313,7 @@ Now run it (two threads should be enough to trigger the data race checks): OMP_NUM_THREADS=2 ./pi_example ``` -However, contrary to the expection, TSan report data races, even though the +However, contrary to the expectation, TSan reports data races, even though the result is always computed correctly: ``` @@ -427,7 +427,7 @@ Now, adjust the command line as requested: OMP_NUM_THREADS=1 ARCHER_OPTIONS='verbose=1' TSAN_OPTIONS='ignore_noninstrumented_modules=1' ./pi_example ``` -and we gete: +and we get: ``` Archer detected OpenMP application with TSan, supplying OpenMP synchronization semantics @@ -505,7 +505,7 @@ This is what we work on next. ## Build the stage1 demo using a custom toolchain -from: https://bazel.build/tutorials/ccp-toolchain-config +from: https://bazel.build/tutorials/cpp-toolchain-config This uses a system-provided `clang` installation (instead of the default compiler Bazel uses). @@ -529,7 +529,7 @@ bazel build --config=clang_config //abseil-hello:hello_main bazel-bin/abseil-hello/hello_main "from Abseil" ``` -Fixed by added `-lm` to standard linker flags. This was inspired by: +Fixed by adding `-lm` to standard linker flags. This was inspired by: https://github.com/bazelbuild/bazel/issues/934#issuecomment-193474914 Even the unit test works if a recent version of `googletest` is used: @@ -633,7 +633,7 @@ Delete the Bazel cache: bazel clean --expunge ``` -The PDF files of the slides mentioned in this articles are mirrored locally +The PDF files of the slides mentioned in this article are mirrored locally here: https://drive.google.com/drive/folders/1NBNTr4jDQy951CoG-AYqpDfKKpNKNSzh?usp=sharing From 4f40e831047088ffa388ab97715af9de8f6b385e Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Thu, 9 Jul 2026 22:28:58 +0200 Subject: [PATCH 07/12] Restore external optimizer behavioral tests (#623) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run every solver once on explicit Solov'ev and CTH-like inputs. Keep the strict internal-state comparison on the reproducible 2D case, and verify force balance and energy in 3D without optional diagnostics or repeated evaluations. Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com> --- tests/test_external_optimizers.py | 107 ++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 tests/test_external_optimizers.py diff --git a/tests/test_external_optimizers.py b/tests/test_external_optimizers.py new file mode 100644 index 000000000..5e20b7fa0 --- /dev/null +++ b/tests/test_external_optimizers.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""External optimizers reach force balance on axisymmetric and 3D cases. + +The Solov'ev equilibrium has a reproducible internal state, so the 2D test compares the +full state vector with the native solver. In 3D, the poloidal parameterization is not +unique and spectral condensation only regularizes that coordinate freedom. The 3D test +therefore compares the residual and energy instead of coordinate-dependent Fourier +coefficients. +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from external_optimizers import ( # type: ignore + make_model, + reference_equilibrium, + residual, + solve_newton_hvp, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_preconditioned_descent, +) + +ROOT = Path(__file__).resolve().parents[1] +SOLOVEV = ROOT / "examples" / "data" / "solovev.json" +CTH_LIKE = ROOT / "examples" / "data" / "cth_like_fixed_bdy.json" +CMA = ROOT / "src" / "vmecpp" / "cpp" / "vmecpp" / "test_data" / "cma.json" + +SOLVERS = ( + solve_preconditioned_descent, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_newton_hvp, +) + + +def _solve_all(input_path): + return {solver.__name__: solver(input_path, ns=11) for solver in SOLVERS} + + +@pytest.fixture(scope="module") +def reference_2d(): + return reference_equilibrium(SOLOVEV, ns=11) + + +@pytest.fixture(scope="module") +def reference_3d(): + return reference_equilibrium(CTH_LIKE, ns=11) + + +@pytest.fixture(scope="module") +def solutions_2d(): + return _solve_all(SOLOVEV) + + +@pytest.fixture(scope="module") +def solutions_3d(): + return _solve_all(CTH_LIKE) + + +def test_optimizers_reach_2d_equilibrium(reference_2d, solutions_2d): + x_star, w_star = reference_2d + for name, (x, result) in solutions_2d.items(): + assert result.residual_norm < 1e-7, name + assert abs(result.energy - w_star) < 1e-8, name + assert np.linalg.norm(x - x_star) < 1e-5, name + + +def test_optimizers_reach_3d_force_balance(reference_3d, solutions_3d): + _, w_star = reference_3d + for name, (_, result) in solutions_3d.items(): + assert result.residual_norm < 1e-7, name + assert np.isclose(result.energy, w_star, rtol=1e-4, atol=0.0), name + + +def test_preconditioner_reduces_newton_krylov_force_evaluations(solutions_3d): + plain = solutions_3d[solve_newton_krylov.__name__][1] + preconditioned = solutions_3d[solve_newton_krylov_preconditioned.__name__][1] + assert preconditioned.force_evals < plain.force_evals + + +def test_hvp_newton_uses_fewer_outer_iterations_than_descent(solutions_3d): + newton = solutions_3d[solve_newton_hvp.__name__][1] + descent = solutions_3d[solve_preconditioned_descent.__name__][1] + assert newton.outer_iters < 20 + assert newton.outer_iters < descent.outer_iters + + +def test_cma_cold_start_exercises_non_axisymmetric_paths(): + model = make_model(CMA, ns=25) + x0 = np.asarray(model.get_state(), float) + f0 = residual(model)(x0) + assert np.all(np.isfinite(f0)) + assert np.linalg.norm(f0) > 0.0 + + rng = np.random.default_rng(0) + v = rng.standard_normal(x0.size) + hv = np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float) + assert np.all(np.isfinite(hv)) + assert np.linalg.norm(hv) > 0.0 From 96f3cb393741abecef5428725a64347c3507ffe9 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Sun, 14 Jun 2026 10:08:04 +0200 Subject: [PATCH 08/12] examples: adjoint boundary sensitivities; SIMSOPT analytic gradient Add the implicit-function adjoint that turns VMEC++ into a gradient-providing equilibrium component for SIMSOPT, the original goal. vmecpp_adjoint.py: for a converged fixed-boundary equilibrium F_I(x)=0, the boundary sensitivity of a scalar objective J follows from H_II lambda = dJ/dx_I, dJ/dx_B = dJ/dx_B - (dF_I/dx_B)^T lambda, with H the symmetric Hessian of the augmented functional. It is matrix-free via hessian_vector_product and apply_preconditioner (the SPD interior system is solved with preconditioned CG). One Hessian solve gives the whole boundary gradient, versus one equilibrium re-solve per boundary DOF for finite differences. simsopt_vmec_gradient.py: VmecEnergy wraps this as a SIMSOPT Optimizable whose dJ is the adjoint gradient, plus a gradient-cost benchmark. Verified: the adjoint gradient matches brute-force re-solve finite differences (rel 2.4e-4) and the SIMSOPT Optimizable's dJ matches finite differences of J (rel ~1e-6). On solovev (ns=11, 18 boundary DOFs) the adjoint boundary gradient costs 762 force evaluations versus 9112 for finite differences (12x), and the gap grows with the boundary DOF count. --- examples/simsopt_vmec_gradient.py | 154 +++++++++++++++++++++++++++ examples/vmecpp_adjoint.py | 169 ++++++++++++++++++++++++++++++ tests/test_adjoint.py | 51 +++++++++ tests/test_simsopt_gradient.py | 57 ++++++++++ 4 files changed, 431 insertions(+) create mode 100644 examples/simsopt_vmec_gradient.py create mode 100644 examples/vmecpp_adjoint.py create mode 100644 tests/test_adjoint.py create mode 100644 tests/test_simsopt_gradient.py diff --git a/examples/simsopt_vmec_gradient.py b/examples/simsopt_vmec_gradient.py new file mode 100644 index 000000000..2240d2a66 --- /dev/null +++ b/examples/simsopt_vmec_gradient.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Use VMEC++ as a gradient-providing equilibrium component for SIMSOPT. + +A SIMSOPT optimization over a plasma boundary normally differentiates VMEC by +finite differences: one equilibrium re-solve per boundary degree of freedom and +per outer iteration. VMEC++ instead provides the boundary gradient analytically +through the implicit-function adjoint (``vmecpp_adjoint.boundary_gradient``): one +extra Hessian solve, independent of the number of boundary DOFs. + +``VmecEnergy`` wraps that as a SIMSOPT ``Optimizable`` whose objective is the MHD +energy of the converged equilibrium and whose ``dJ`` is the adjoint gradient. +``optimize_to_target`` runs a gradient-based optimization of the boundary toward +a target energy, with the analytic gradient or with finite differences, and +reports the cost (forward-model evaluations counted inside VMEC++, outer +iterations, wall time) so the two can be compared. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from vmecpp_adjoint import ( + DEFAULT_INPUT, + boundary_gradient, + finite_difference_boundary_gradient, + make_model, + mhd_energy, + partition, + solve_interior, +) + + +class VmecBoundaryProblem: + """Equilibrium energy and its boundary gradient, cached per boundary state.""" + + def __init__(self, input_path: Path = DEFAULT_INPUT, ns: int = 11): + self.model = make_model(input_path, ns) + self.model.solve() + self.ns = ns + self.interior, self.boundary = partition(self.model, ns) + self._x_full = np.asarray(self.model.get_state(), float).copy() + self._cached_p = self._x_full[self.boundary].copy() + + @property + def x0(self): + return self._x_full[self.boundary].copy() + + def _resolve(self, p): + p = np.asarray(p, float) + if not np.array_equal(p, self._cached_p): + self._x_full = solve_interior( + self.model, self._x_full, self.interior, self.boundary, p + ) + self._cached_p = p.copy() + + def value(self, p): + self._resolve(p) + self.model.set_state(np.ascontiguousarray(self._x_full)) + self.model.evaluate(2, 2, False) + return self.model.mhd_energy + + def gradient(self, p): + self._resolve(p) + return boundary_gradient( + self.model, self._x_full, self.interior, self.boundary, mhd_energy + ) + + +def make_simsopt_optimizable(problem: VmecBoundaryProblem): + """Wrap the problem as a SIMSOPT Optimizable exposing an analytic gradient.""" + # Imported lazily so the rest of the module (and the gradient benchmark) work + # without SIMSOPT installed. + from simsopt._core import Optimizable # noqa: PLC0415 + from simsopt._core.derivative import Derivative, derivative_dec # noqa: PLC0415 + + class VmecEnergy(Optimizable): + def __init__(self): + x0 = problem.x0 + super().__init__(x0=x0, names=[f"boundary{i}" for i in range(x0.size)]) + + def J(self): + return problem.value(self.local_full_x) + + @derivative_dec + def dJ(self): + return Derivative({self: problem.gradient(self.local_full_x)}) + + return VmecEnergy() + + +@dataclass +class GradResult: + method: str + force_evals: int + seconds: float + gradient: np.ndarray + + +def gradient_cost(input_path: Path = DEFAULT_INPUT, ns: int = 11, analytic=True): + """Cost of one full boundary gradient at the converged equilibrium. + + This is what an external optimizer pays per iteration. The analytic adjoint + needs one Hessian solve regardless of the number of boundary DOFs; finite + differences re-converge the equilibrium twice per boundary DOF. + """ + problem = VmecBoundaryProblem(input_path, ns) + x_star = problem._x_full.copy() + interior, boundary = problem.interior, problem.boundary + problem.model.reset_force_eval_count() + t0 = time.perf_counter() + if analytic: + g_dict = None + g = boundary_gradient(problem.model, x_star, interior, boundary, mhd_energy) + else: + g_dict = finite_difference_boundary_gradient( + problem.model, + x_star, + interior, + boundary, + mhd_energy, + range(boundary.size), + ) + g = np.array([g_dict[j] for j in range(boundary.size)]) + return GradResult( + "analytic adjoint" if analytic else "finite differences", + problem.model.force_eval_count, + time.perf_counter() - t0, + g, + ) + + +def main(): + analytic = gradient_cost(analytic=True) + fd = gradient_cost(analytic=False) + rel = np.linalg.norm(analytic.gradient - fd.gradient) / np.linalg.norm(fd.gradient) + n_boundary = analytic.gradient.size + print(f"boundary gradient cost ({n_boundary} boundary DOFs, solovev ns=11)\n") + print(f"{'method':20s} {'F-evals':>9s} {'time[s]':>8s}") + for r in (fd, analytic): + print(f"{r.method:20s} {r.force_evals:9d} {r.seconds:8.2f}") + print( + f"\nspeedup (force evals): {fd.force_evals / max(analytic.force_evals, 1):.1f}x" + f" gradient agreement: {rel:.1e}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/vmecpp_adjoint.py b/examples/vmecpp_adjoint.py new file mode 100644 index 000000000..e4b27f3bf --- /dev/null +++ b/examples/vmecpp_adjoint.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Adjoint sensitivity of a converged VMEC++ equilibrium to its boundary. + +A fixed-boundary equilibrium satisfies the interior force balance F_I(x) = 0, +where x is the decomposed internal-basis state and F is the gradient of VMEC's +augmented functional. The outermost flux surface (the boundary) is the last +radial block of the state and is held fixed during the solve. For a scalar +objective J(x), the sensitivity to the boundary degrees of freedom follows from +the implicit function theorem: + + dJ/dx_B = dJ/dx_B|_x - (dF_I/dx_B)^T lambda, H_II lambda = dJ/dx_I, + +with H = dF/dx the (symmetric) Hessian of the augmented functional. Every +operator is matrix-free and already exposed by VmecModel: the Hessian-vector +product (``hessian_vector_product``) and the preconditioner +(``apply_preconditioner``), used to solve the adjoint system. Only one Hessian +solve is needed for the full boundary gradient, versus one equilibrium re-solve +per boundary degree of freedom for finite differences. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from scipy.sparse.linalg import LinearOperator, cg + +try: + from vmecpp.cpp import _vmecpp +except ImportError: + import _vmecpp + +DEFAULT_INPUT = ( + Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" +) + + +def make_model(input_path: Path = DEFAULT_INPUT, ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(input_path)), ns) + + +def partition(model, ns: int): + """Indices of the interior (free) and boundary (LCFS) state components.""" + k = model.mpol * (model.ntor + 1) + n = np.asarray(model.get_state()).size + per_span = ns * k + n_span = n // per_span + boundary = [] + for s in range(n_span): + boundary.extend(range(s * per_span + (ns - 1) * k, s * per_span + ns * k)) + boundary = np.array(sorted(boundary)) + interior = np.setdiff1d(np.arange(n), boundary) + return interior, boundary + + +def _raw_force(model, x): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), float) + + +def _interior_operators(model, x, interior): + # The caller must set the base state to x and assemble the preconditioner + # (evaluate(2, 2, True)) before using these. hessian_vector_product uses the + # current state as its base point and restores it, so no per-matvec state + # update is needed: that keeps each Hessian matvec at two force evaluations. + n = x.size + ni = interior.size + + def hii(vi): + v = np.zeros(n) + v[interior] = vi + return np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float)[ + interior + ] + + def mii(bi): + v = np.zeros(n) + v[interior] = bi + return np.asarray(model.apply_preconditioner(np.ascontiguousarray(v)), float)[ + interior + ] + + return ( + LinearOperator((ni, ni), matvec=hii), + LinearOperator((ni, ni), matvec=mii), + ) + + +def solve_interior(model, x0, interior, boundary, x_boundary, tol=1e-10, max_newton=40): + """Converge the interior to force balance with the boundary held fixed.""" + x = np.asarray(x0, float).copy() + x[boundary] = x_boundary + for _ in range(max_newton): + f = _raw_force(model, x) + if np.linalg.norm(f[interior]) < tol: + break + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, True) # assemble preconditioner + set base state + h_op, m_op = _interior_operators(model, x, interior) + dxi, _ = cg(h_op, -f[interior], M=m_op, rtol=1e-4, maxiter=200) + x[interior] += dxi + return x + + +def objective_state_gradient(model, x, objective, h=1e-6): + """Partial derivative dJ/dx at fixed state, by central finite differences.""" + n = x.size + g = np.zeros(n) + for i in range(n): + xp = x.copy() + xp[i] += h + model.set_state(np.ascontiguousarray(xp)) + model.evaluate(2, 2, False) + jp = objective(model) + xm = x.copy() + xm[i] -= h + model.set_state(np.ascontiguousarray(xm)) + model.evaluate(2, 2, False) + jm = objective(model) + g[i] = (jp - jm) / (2 * h) + return g + + +def boundary_gradient(model, x_star, interior, boundary, objective, h=1e-6): + """Adjoint gradient dJ/dx_B at the converged equilibrium x_star.""" + n = x_star.size + dj = objective_state_gradient(model, x_star, objective, h) + model.set_state(np.ascontiguousarray(x_star)) + model.evaluate(2, 2, True) # assemble preconditioner + set base state to x_star + h_op, m_op = _interior_operators(model, x_star, interior) + lam, _ = cg(h_op, dj[interior], M=m_op, rtol=1e-7, maxiter=400) + embedded = np.zeros(n) + embedded[interior] = lam + model.set_state(np.ascontiguousarray(x_star)) + model.evaluate(2, 2, False) + coupling = np.asarray( + model.hessian_vector_product(np.ascontiguousarray(embedded)), float + )[boundary] + return dj[boundary] - coupling + + +def finite_difference_boundary_gradient( + model, x_star, interior, boundary, objective, dofs, h=1e-5 +): + """Reference gradient: re-solve the interior for each perturbed boundary DOF.""" + g = {} + for j in dofs: + xbp = x_star[boundary].copy() + xbp[j] += h + xp = solve_interior(model, x_star, interior, boundary, xbp) + model.set_state(np.ascontiguousarray(xp)) + model.evaluate(2, 2, False) + jp = objective(model) + xbm = x_star[boundary].copy() + xbm[j] -= h + xm = solve_interior(model, x_star, interior, boundary, xbm) + model.set_state(np.ascontiguousarray(xm)) + model.evaluate(2, 2, False) + jm = objective(model) + g[j] = (jp - jm) / (2 * h) + return g + + +def mhd_energy(model): + return model.mhd_energy diff --git a/tests/test_adjoint.py b/tests/test_adjoint.py new file mode 100644 index 000000000..0d740d6be --- /dev/null +++ b/tests/test_adjoint.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""The adjoint boundary gradient matches brute-force finite differences. + +dJ/d(boundary) from one Hessian solve (implicit-function adjoint) agrees with the +reference gradient obtained by re-converging the interior equilibrium for each +perturbed boundary degree of freedom. J here is the MHD energy of the converged +equilibrium. +""" + +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from vmecpp_adjoint import ( + boundary_gradient, + finite_difference_boundary_gradient, + make_model, + mhd_energy, + partition, +) + + +def test_adjoint_matches_finite_difference(): + ns = 11 + model = make_model(ns=ns) + model.solve() + x_star = np.asarray(model.get_state(), float).copy() + interior, boundary = partition(model, ns) + + g_adjoint = boundary_gradient(model, x_star, interior, boundary, mhd_energy) + + # Reference on a representative subset (each requires interior re-solves). + dofs = [0, 2, 9] + g_fd = finite_difference_boundary_gradient( + model, x_star, interior, boundary, mhd_energy, dofs + ) + + scale = max(np.linalg.norm(g_adjoint), 1e-30) + for j in dofs: + assert abs(g_adjoint[j] - g_fd[j]) < 1e-3 * scale + + +if __name__ == "__main__": + import pytest + + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_simsopt_gradient.py b/tests/test_simsopt_gradient.py new file mode 100644 index 000000000..4080079fe --- /dev/null +++ b/tests/test_simsopt_gradient.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VMEC++ exposes an analytic boundary gradient to SIMSOPT. + +The VmecEnergy Optimizable's analytic dJ (the implicit-function adjoint) matches +finite differences of its objective, and computing it is much cheaper than the +conventional finite-difference boundary gradient (which re-solves the +equilibrium per boundary degree of freedom). +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) + +pytest.importorskip("simsopt") + +from simsopt_vmec_gradient import ( + VmecBoundaryProblem, + gradient_cost, + make_simsopt_optimizable, +) + + +def test_simsopt_optimizable_gradient_matches_fd(): + problem = VmecBoundaryProblem(ns=11) + opt = make_simsopt_optimizable(problem) + g = np.asarray(opt.dJ(), float) + + p0 = np.asarray(opt.local_full_x, float) + h = 1e-5 + scale = max(np.linalg.norm(g), 1e-30) + for j in (0, 2, 9): + pp = p0.copy() + pp[j] += h + opt.local_full_x = pp + jp = opt.J() + pm = p0.copy() + pm[j] -= h + opt.local_full_x = pm + jm = opt.J() + opt.local_full_x = p0 + assert abs(g[j] - (jp - jm) / (2 * h)) < 1e-3 * scale + + +def test_adjoint_gradient_cheaper_than_finite_difference(): + analytic = gradient_cost(analytic=True) + fd = gradient_cost(analytic=False) + # Same gradient, far fewer force evaluations (advantage grows with #DOFs). + rel = np.linalg.norm(analytic.gradient - fd.gradient) / np.linalg.norm(fd.gradient) + assert rel < 1e-2 + assert analytic.force_evals < fd.force_evals From d6d109127f5919196c0a822fd6c7578a1e93c715 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Sun, 14 Jun 2026 10:52:21 +0200 Subject: [PATCH 09/12] examples: fix adjoint on 3D (GMRES for indefinite Hessian) + globalize Two correctness fixes for stiff 3D equilibria (cth_like): - VMEC's augmented-Lagrangian Hessian is symmetric *indefinite* (the lambda constraint makes it a saddle, not a minimum), so CG silently gives the wrong adjoint there. Use GMRES, which handles indefinite systems, for the H_II solve and the interior Newton solve. With a loose, restarted tolerance the adjoint solve stays cheap. - Add a backtracking line search to solve_interior so the interior re-solve (used by the SIMSOPT wrapper and the finite-difference reference) converges on 3D instead of overshooting. Verified with a directional-derivative check against a re-converged finite-difference reference: solovev 1.5e-4, cth_like 2.2e-2 relative; both previously agreed only in 2D. Boundary-gradient cost on solovev: 626 force evaluations (analytic adjoint) versus 10460 (finite differences). --- examples/vmecpp_adjoint.py | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/examples/vmecpp_adjoint.py b/examples/vmecpp_adjoint.py index e4b27f3bf..a85ea2006 100644 --- a/examples/vmecpp_adjoint.py +++ b/examples/vmecpp_adjoint.py @@ -26,7 +26,7 @@ from pathlib import Path import numpy as np -from scipy.sparse.linalg import LinearOperator, cg +from scipy.sparse.linalg import LinearOperator, gmres try: from vmecpp.cpp import _vmecpp @@ -90,19 +90,34 @@ def mii(bi): ) -def solve_interior(model, x0, interior, boundary, x_boundary, tol=1e-10, max_newton=40): - """Converge the interior to force balance with the boundary held fixed.""" +def solve_interior(model, x0, interior, boundary, x_boundary, tol=1e-10, max_newton=80): + """Converge the interior to force balance with the boundary held fixed. + + Preconditioned Newton-Krylov on the interior residual with a backtracking + line search; the line search is required for stiff 3D equilibria, where the + full Newton step overshoots. + """ x = np.asarray(x0, float).copy() x[boundary] = x_boundary for _ in range(max_newton): f = _raw_force(model, x) - if np.linalg.norm(f[interior]) < tol: + norm0 = np.linalg.norm(f[interior]) + if norm0 < tol: break model.set_state(np.ascontiguousarray(x)) model.evaluate(2, 2, True) # assemble preconditioner + set base state h_op, m_op = _interior_operators(model, x, interior) - dxi, _ = cg(h_op, -f[interior], M=m_op, rtol=1e-4, maxiter=200) - x[interior] += dxi + dxi, _ = gmres(h_op, -f[interior], M=m_op, rtol=1e-4, maxiter=300) + alpha = 1.0 + for _ in range(30): + xt = x.copy() + xt[interior] += alpha * dxi + if np.linalg.norm(_raw_force(model, xt)[interior]) < norm0: + break + alpha *= 0.5 + else: + break # no decrease found; stop + x[interior] += alpha * dxi return x @@ -132,7 +147,7 @@ def boundary_gradient(model, x_star, interior, boundary, objective, h=1e-6): model.set_state(np.ascontiguousarray(x_star)) model.evaluate(2, 2, True) # assemble preconditioner + set base state to x_star h_op, m_op = _interior_operators(model, x_star, interior) - lam, _ = cg(h_op, dj[interior], M=m_op, rtol=1e-7, maxiter=400) + lam, _ = gmres(h_op, dj[interior], M=m_op, rtol=1e-6, restart=100, maxiter=30) embedded = np.zeros(n) embedded[interior] = lam model.set_state(np.ascontiguousarray(x_star)) From 2f258ed37acfbf829a6b99c36fcaeae0a67fdd11 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Thu, 9 Jul 2026 22:30:32 +0200 Subject: [PATCH 10/12] Fix adjoint checks on current main --- examples/simsopt_vmec_gradient.py | 15 ++++++++++----- examples/vmecpp_adjoint.py | 15 ++++++--------- tests/test_adjoint.py | 7 +++---- tests/test_simsopt_gradient.py | 10 +++++----- 4 files changed, 24 insertions(+), 23 deletions(-) diff --git a/examples/simsopt_vmec_gradient.py b/examples/simsopt_vmec_gradient.py index 2240d2a66..54f188789 100644 --- a/examples/simsopt_vmec_gradient.py +++ b/examples/simsopt_vmec_gradient.py @@ -77,7 +77,11 @@ def make_simsopt_optimizable(problem: VmecBoundaryProblem): # Imported lazily so the rest of the module (and the gradient benchmark) work # without SIMSOPT installed. from simsopt._core import Optimizable # noqa: PLC0415 - from simsopt._core.derivative import Derivative, derivative_dec # noqa: PLC0415 + from simsopt._core.derivative import ( # noqa: PLC0415 + Derivative, + OptimizableDefaultDict, + derivative_dec, + ) class VmecEnergy(Optimizable): def __init__(self): @@ -89,7 +93,8 @@ def J(self): @derivative_dec def dJ(self): - return Derivative({self: problem.gradient(self.local_full_x)}) + data = OptimizableDefaultDict({self: problem.gradient(self.local_full_x)}) + return Derivative(data) return VmecEnergy() @@ -105,9 +110,9 @@ class GradResult: def gradient_cost(input_path: Path = DEFAULT_INPUT, ns: int = 11, analytic=True): """Cost of one full boundary gradient at the converged equilibrium. - This is what an external optimizer pays per iteration. The analytic adjoint - needs one Hessian solve regardless of the number of boundary DOFs; finite - differences re-converge the equilibrium twice per boundary DOF. + This is what an external optimizer pays per iteration. The analytic adjoint needs + one Hessian solve regardless of the number of boundary DOFs; finite differences re- + converge the equilibrium twice per boundary DOF. """ problem = VmecBoundaryProblem(input_path, ns) x_star = problem._x_full.copy() diff --git a/examples/vmecpp_adjoint.py b/examples/vmecpp_adjoint.py index a85ea2006..f10b63d12 100644 --- a/examples/vmecpp_adjoint.py +++ b/examples/vmecpp_adjoint.py @@ -28,10 +28,7 @@ import numpy as np from scipy.sparse.linalg import LinearOperator, gmres -try: - from vmecpp.cpp import _vmecpp -except ImportError: - import _vmecpp +from vmecpp.cpp import _vmecpp # type: ignore DEFAULT_INPUT = ( Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" @@ -85,17 +82,17 @@ def mii(bi): ] return ( - LinearOperator((ni, ni), matvec=hii), - LinearOperator((ni, ni), matvec=mii), + LinearOperator((ni, ni), matvec=hii), # type: ignore[call-overload] + LinearOperator((ni, ni), matvec=mii), # type: ignore[call-overload] ) def solve_interior(model, x0, interior, boundary, x_boundary, tol=1e-10, max_newton=80): """Converge the interior to force balance with the boundary held fixed. - Preconditioned Newton-Krylov on the interior residual with a backtracking - line search; the line search is required for stiff 3D equilibria, where the - full Newton step overshoots. + Preconditioned Newton-Krylov on the interior residual with a backtracking line + search; the line search is required for stiff 3D equilibria, where the full Newton + step overshoots. """ x = np.asarray(x0, float).copy() x[boundary] = x_boundary diff --git a/tests/test_adjoint.py b/tests/test_adjoint.py index 0d740d6be..54a45c1a1 100644 --- a/tests/test_adjoint.py +++ b/tests/test_adjoint.py @@ -5,9 +5,8 @@ """The adjoint boundary gradient matches brute-force finite differences. dJ/d(boundary) from one Hessian solve (implicit-function adjoint) agrees with the -reference gradient obtained by re-converging the interior equilibrium for each -perturbed boundary degree of freedom. J here is the MHD energy of the converged -equilibrium. +reference gradient obtained by re-converging the interior equilibrium for each perturbed +boundary degree of freedom. J here is the MHD energy of the converged equilibrium. """ import sys @@ -16,7 +15,7 @@ import numpy as np sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) -from vmecpp_adjoint import ( +from vmecpp_adjoint import ( # type: ignore boundary_gradient, finite_difference_boundary_gradient, make_model, diff --git a/tests/test_simsopt_gradient.py b/tests/test_simsopt_gradient.py index 4080079fe..90f4f97c5 100644 --- a/tests/test_simsopt_gradient.py +++ b/tests/test_simsopt_gradient.py @@ -4,10 +4,10 @@ # SPDX-License-Identifier: MIT """VMEC++ exposes an analytic boundary gradient to SIMSOPT. -The VmecEnergy Optimizable's analytic dJ (the implicit-function adjoint) matches -finite differences of its objective, and computing it is much cheaper than the -conventional finite-difference boundary gradient (which re-solves the -equilibrium per boundary degree of freedom). +The VmecEnergy Optimizable's analytic dJ (the implicit-function adjoint) matches finite +differences of its objective, and computing it is much cheaper than the conventional +finite-difference boundary gradient (which re-solves the equilibrium per boundary degree +of freedom). """ import sys @@ -20,7 +20,7 @@ pytest.importorskip("simsopt") -from simsopt_vmec_gradient import ( +from simsopt_vmec_gradient import ( # type: ignore VmecBoundaryProblem, gradient_cost, make_simsopt_optimizable, From d0ef637ff6e7aaf009bc83143d7a327481bd4316 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Thu, 9 Jul 2026 23:27:27 +0200 Subject: [PATCH 11/12] Improve docstring formatting in simsopt_vmec_gradient.py Fix formatting in docstring for clarity. --- examples/simsopt_vmec_gradient.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/simsopt_vmec_gradient.py b/examples/simsopt_vmec_gradient.py index 54f188789..78eafde6e 100644 --- a/examples/simsopt_vmec_gradient.py +++ b/examples/simsopt_vmec_gradient.py @@ -13,7 +13,7 @@ ``VmecEnergy`` wraps that as a SIMSOPT ``Optimizable`` whose objective is the MHD energy of the converged equilibrium and whose ``dJ`` is the adjoint gradient. ``optimize_to_target`` runs a gradient-based optimization of the boundary toward -a target energy, with the analytic gradient or with finite differences, and +a target energy, with analytic gradient or with finite differences, and reports the cost (forward-model evaluations counted inside VMEC++, outer iterations, wall time) so the two can be compared. """ From f0cb35a25d825fafe795dc8db03d59500e6fe134 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Fri, 10 Jul 2026 10:16:57 +0200 Subject: [PATCH 12/12] Make raw VMEC forces history independent (#626) * Make raw VMEC forces history independent Separate the legacy previous-residual m=1 projection used by native iteration from the exact constrained projection used by external force evaluations. Cover both directions across the residual threshold. * Use compact storage for constraint policy * Name the m=1 gauge policy explicitly --- src/vmecpp/_iteration.py | 2 +- .../vmec/ideal_mhd_model/ideal_mhd_model.cc | 7 +++- .../vmec/ideal_mhd_model/ideal_mhd_model.h | 3 +- .../cpp/vmecpp/vmec/pybind11/pybind_vmec.cc | 11 +++-- tests/test_internal_gradient.py | 42 +++++++++++++++++++ 5 files changed, 58 insertions(+), 7 deletions(-) diff --git a/src/vmecpp/_iteration.py b/src/vmecpp/_iteration.py index b019a79fa..e54a8cfc5 100644 --- a/src/vmecpp/_iteration.py +++ b/src/vmecpp/_iteration.py @@ -259,7 +259,7 @@ def solve_equilibrium( iter2 = force_iteration - bad_resets # Evolve: forward model, then convergence / damping / time step. - model.evaluate(iter1, iter2) + model.evaluate(iter1, iter2, precondition=True, always_fix_m1_gauge=False) fsqr, fsqz, fsql = model.fsqr, model.fsqz, model.fsql rr = model.restart_reason finite = math.isfinite(fsqr) and math.isfinite(fsqz) and math.isfinite(fsql) diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc index 1ec40ce51..66b746e1c 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc @@ -421,7 +421,8 @@ absl::StatusOr IdealMhdModel::update( bool& m_need_restart, int& m_last_preconditioner_update, int& m_last_full_update_nestor, FlowControl& m_fc, const int iter1, const int iter2, const VmecCheckpoint& checkpoint, - const int iterations_before_checkpointing, bool verbose) { + const int iterations_before_checkpointing, bool verbose, + bool always_fix_m1_gauge) { // An axis re-guess after a bad Jacobian can repopulate high geometry modes // directly, bypassing the force mask; clear them on the state each iteration. if (s_.mpolGeometry < s_.mpol || s_.ntorGeometry < s_.ntor) { @@ -813,7 +814,9 @@ absl::StatusOr IdealMhdModel::update( m_decomposed_f.m1Constraint(1.0 / std::numbers::sqrt2); // v8.50: ADD iter2<2 so reset= works - if (m_fc.fsqz < 1.0e-6 || iter2 < 2) { + const bool fix_m1_gauge = + always_fix_m1_gauge || m_fc.fsqz < 1.0e-6 || iter2 < 2; + if (fix_m1_gauge) { // ensure that the m=1 constraint is satisfied exactly // --> the corresponding m=1 coeffs of R,Z are constrained to be zero // and thus must not be "forced" (by the time evol using gc) away from diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h index b753e53f5..da3cf19fc 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h @@ -70,7 +70,8 @@ class IdealMhdModel { bool& m_need_restart, int& m_last_preconditioner_update, int& m_last_full_update_nestor, FlowControl& m_fc, const int iter1, const int iter2, const VmecCheckpoint& checkpoint = VmecCheckpoint::NONE, - const int iterations_before_checkpointing = INT_MAX, bool verbose = true); + const int iterations_before_checkpointing = INT_MAX, bool verbose = true, + bool always_fix_m1_gauge = false); // Coordinates which inverse-DFT routine to call for computing // the flux surface geometry and lambda on it from the provided Fourier diff --git a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc index cda5d0ad9..5e2510d7d 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc @@ -226,7 +226,11 @@ class VmecModel { // lambda-constraint components. That raw gradient is what gradient-based // optimizers minimizing the MHD energy functional need; mhd_energy is already // set earlier in update(), so it is valid at the checkpoint too. - void Evaluate(int iter1, int iter2, bool precondition = true) { + // The native iteration leaves the m=1 gauge free until the previous Z + // residual crosses its threshold. External evaluations fix it immediately so + // F(x) does not depend on the previously evaluated state. + void Evaluate(int iter1, int iter2, bool precondition = true, + bool always_fix_m1_gauge = true) { ++force_eval_count_; bool need_restart = false; std::string error_message; @@ -255,7 +259,7 @@ class VmecModel { *vmec_->decomposed_f_[0], *vmec_->physical_f_[0], need_restart, last_preconditioner_update_, last_full_update_nestor_, vmec_->fc_, iter1, iter2, checkpoint, checkpoint_after, - /*verbose=*/false); + /*verbose=*/false, always_fix_m1_gauge); if (!s.ok()) { error_message = std::string(s.status().message()); } @@ -1260,7 +1264,8 @@ PYBIND11_MODULE(_vmecpp, m) { .def_static("create", &VmecModel::Create, py::arg("indata"), py::arg("ns"), py::arg("initial_state") = std::nullopt) .def("evaluate", &VmecModel::Evaluate, py::arg("iter1"), py::arg("iter2"), - py::arg("precondition") = true) + py::arg("precondition") = true, + py::arg("always_fix_m1_gauge") = true) .def_property_readonly("need_restart", &VmecModel::need_restart) .def("perform_time_step", &VmecModel::PerformTimeStep, py::arg("velocity_scale"), py::arg("conjugation_parameter"), diff --git a/tests/test_internal_gradient.py b/tests/test_internal_gradient.py index 17d36f65f..16d26825e 100644 --- a/tests/test_internal_gradient.py +++ b/tests/test_internal_gradient.py @@ -17,10 +17,17 @@ from pathlib import Path import numpy as np +import pytest from vmecpp.cpp import _vmecpp # type: ignore SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" +CTH_LIKE = ( + Path(__file__).resolve().parents[1] + / "examples" + / "data" + / "cth_like_fixed_bdy.json" +) def _model(ns: int = 11): @@ -63,3 +70,38 @@ def test_cold_start_is_excluded(): m = _model() m.evaluate(2, 2, False) assert np.all(np.isfinite(np.asarray(m.get_forces(), float))) + + +@pytest.fixture(scope="module") +def force_histories(): + indata = _vmecpp.VmecINDATA.from_file(str(CTH_LIKE)) + reference = _vmecpp.VmecModel.create(indata, 11) + reference.solve() + low_residual_state = np.asarray(reference.get_state(), float).copy() + + model = _vmecpp.VmecModel.create(indata, 11) + high_residual_state = np.asarray(model.get_state(), float).copy() + model.evaluate(2, 2, False) + high_after_high = np.asarray(model.get_forces(), float).copy() + + model.set_state(np.ascontiguousarray(low_residual_state)) + model.evaluate(2, 2, False) + low_after_high = np.asarray(model.get_forces(), float).copy() + model.evaluate(2, 2, False) + low_after_low = np.asarray(model.get_forces(), float).copy() + + model.set_state(np.ascontiguousarray(high_residual_state)) + model.evaluate(2, 2, False) + high_after_low = np.asarray(model.get_forces(), float).copy() + + return high_after_high, high_after_low, low_after_high, low_after_low + + +def test_raw_force_is_independent_of_high_residual_history(force_histories): + _, _, low_after_high, low_after_low = force_histories + np.testing.assert_array_equal(low_after_high, low_after_low) + + +def test_raw_force_is_independent_of_low_residual_history(force_histories): + high_after_high, high_after_low, _, _ = force_histories + np.testing.assert_array_equal(high_after_high, high_after_low)