feat: autodiff via Enzyme on LLVM Fortran toolchain

[krystophny]

Motivation

The JAX port pyrabe (https://github.com/itpplasma/pyrabe) is differentiable end-to-end only with respect to Boozer Fourier coefficients, and only on a fixed branch of the well topology: same surface, same field-line label grid, same selected pair of maxima per line, same active global B_max. The VMEC-to-Boozer step in pyrabe wraps the legacy booz_xform extension and is not part of the JAX trace, so the chain

SIMSOPT DOFs -> VMEC -> Boozer -> Lambda_A, Lambda_B, Lambda_S

has no exact gradient today. Inside pyrabe, the discrete pieces (find_two_maxima, jnp.max(line_max_B), the max(1 - B*eta_b, 0) clamp, the rational xi_0 grid) are correct piecewise but jump at topology changes: a minimum-maximum pair born or annihilated, two maxima exchanging which is largest, a split well, an active global maximum switching from one field line to another. Reverse-mode AD through those events gives the gradient of the branch immediately before, not through, the event.

Fortran rabe is the canonical numeric kernel. If we get reverse-mode AD on rabe itself via Enzyme, we obtain:

• exact gradient d Lambda_{A,B,S} / d boozer_modes without a parallel JAX reimplementation that has to be kept in sync;
• a differentiable VMEC-to-Boozer step through the existing boozer_converter Fortran source, closing the SIMSOPT chain to whatever extent boozer_converter itself is differentiable;
• adjoint cost roughly independent of the number of inputs, which is the right cost model for stellarator shape optimization.

The build/llvm-fortran-compilers branch already moves the build toward an LLVM-based Fortran frontend (flang-new / LLVM flang), which is what Enzyme needs. This issue tracks the autodiff work that builds on top.

Concrete request

Enable reverse-mode autodiff of rabe.x (or of a rabe_lib entry point) with respect to a chosen input vector, using Enzyme as the AD pass on the LLVM IR produced by the flang-new build.

Reference inputs to differentiate against, ranked by usefulness:

1. Boozer Fourier coefficients on one surface (matches the pyrabe Boozer-coefficient gradient and lets us cross-check against JAX on the same physics).
2. VMEC wout quantities consumed by src/vmec/boozer_converter.F90, propagating through the Boozer transform.
3. Eventually SIMSOPT-side DOFs, via a chain simsopt -> VMEC -> rabe.

Reference outputs: Lambda_A, Lambda_B, Lambda_S, nu_star_crit per surface, as written to rabe.nc.

Known obstacles in the current sources

These are the spots where Enzyme will either need help or where we should replace the algorithm with a differentiable equivalent. File paths are on main / build/llvm-fortran-compilers.

1. src/utils/find_extrema.f90 does a coarse 1001-point scan followed by repeated subdivision with maxloc and an if ladder. Discrete index selection is not differentiable directly. Replacement: fix the selected extremum by index (treat the coarse argmax as a piecewise-constant choice), then differentiate the Newton refinement via the implicit-function theorem at B_phi(phi*, p) = 0, giving d phi* / d p = -B_{phi,p} / B_{phi,phi}. pyrabe/maxima.py already does this; we want the same trick inside rabe.

2. src/fieldline/make_fieldline.f90:get_global_B_max is a hard maxval over field lines, and calc_nu_star_crit uses a hard minval over B_max(:,1). The gradient flows only through the currently active line and jumps when the active line switches. Decision needed: keep the hard max and accept that the optimizer sees a kink, or expose a smooth surrogate (logsumexp with a documented temperature) behind a switch.

3. src/fieldline/make_fieldline.f90:pick_maximum_on_each_side and the split_maxima warning branch encode a topology change. Enzyme cannot differentiate across the event; we should produce a gradient only when split_maxima == 0, and surface the flag in the output the way the namelist already does, so SIMSOPT-side code can treat it as a constraint.

4. src/utils/integrate.f90 calls QUADPACK dqag, which is adaptive: the subdivision pattern depends on the integrand values. Enzyme can in principle differentiate through it, but adaptive control flow with internal error_stop paths is the hardest case. Two cleaner options:

   • keep dqag for the forward computation, and use the substitution + fixed Gauss-Legendre rule that pyrabe.quadrature uses for the reverse pass on the same integrand; or
   • port the four well integrals (wrapper_lambda_over_B_squared, wrapper_local_radial_drift, wrapper_one_over_B_squared, wrapper_nabla_s_over_B_squared) to a fixed-order Gauss-Legendre rule with the same t**2 = x - a / t**2 = b - x substitution that the Fortran code already applies, then differentiate the fixed rule with Enzyme directly. This matches pyrabe and removes the adaptive branch.

5. error stop on integration or validation failure aborts the process. For AD we need those branches reachable as ordinary error returns or as NaN propagation so the autodiff machinery does not get terminated mid-derivative.

6. src/vmec/boozer_converter.F90 and the libneo field interface: these are the SIMSOPT-side payoff. Differentiability here is a separate workstream; we should first land Enzyme on the rabe-internal path with the Boozer field as input and treat boozer-converter as a later milestone.

Suggested first milestone

One surface, one quantity, fixed maximum topology:

• input: Boozer Fourier coefficients on a fixed s_tor;
• output: Lambda_A for that surface;
• check: Enzyme gradient agrees with central finite differences on rabe.x to ~rel_error_tol_quadkind (currently 1e-6), and agrees with the jax.grad of pyrabe.coefficients.Lambda_A on the same Boozer input to within the documented quadrature tolerance.

Once that holds, extend to Lambda_B, Lambda_S, then to multiple surfaces, then push the input boundary upstream into boozer_converter.

Build prerequisites

• LLVM flang (work in progress on build/llvm-fortran-compilers).
• Enzyme built against the same LLVM version as flang.
• cmake/compiler.cmake already branches on CMAKE_Fortran_COMPILER_ID; add a Flang / LLVMFlang branch with -fno-fortran-main-style flags as needed and a -fpass-plugin=…/LLVMEnzyme-*.so option.
• QUADPACK is BSD-3; either keep as-is for the forward pass or replace the four well integrals as in obstacle 4.

Excluded

• Forward-mode AD: cost model is wrong for shape optimization.
• Differentiable rewrite of libneo field I/O.
• Smoothing the rational xi_0 grid in src/fieldline/fieldline_labels.f90. Freezing this grid per optimization segment is acceptable for a first version.

References

• pyrabe analysis of the same numerics: https://github.com/itpplasma/pyrabe (see src/pyrabe/maxima.py, src/pyrabe/surface_pipeline.py, src/pyrabe/__init__.py for the JAX-side limits this issue is meant to lift on the Fortran side).
• Enzyme: https://enzyme.mit.edu/ and Moses & Churavy, NeurIPS 2020.

[krystophny]

Merging scope from #77 into this issue so #75 remains the canonical rabe tracker for Enzyme AD.

#77 duplicated this issue's main goal, but it added a useful implementation split that should be preserved here:

Fixed-branch design

The Enzyme target should not be the full rabe.x executable. It should be a library-level scalar objective kernel that assumes a fixed branch:

• fixed surface list and field-line grid;
• fixed selected pair of maxima per field line;
• fixed active global B_max for eta_b;
• fixed active minval/maxval entries;
• fixed spline cells/stencils in the VMEC/Boozer layer;
• no disappearing/splitting wells during the differentiated step.

Branch/topology discovery should remain outside the Enzyme-active region. If the branch changes, the optimizer should reject/shrink the step or rerun branch selection; Enzyme will only provide the derivative of the currently executed smooth branch.

Proposed split

Inactive/setup phase:

• read namelist and VMEC/wout files;
• initialize Boozer data through libneo;
• choose surfaces and field-line labels;
• discover/select maxima and detect split_maxima;
• choose active global B_max and other minval/maxval indices;
• choose spline cells/stencils in the field-evaluation layer.

Active Enzyme phase:

• refine already-selected maxima with fixed Newton iterations, or use fixed endpoints for the first milestone;
• evaluate field quantities through fixed-cell/fixed-stencil calls;
• perform fixed Gauss-Legendre quadrature with the square-root endpoint substitution;
• compute field-line integrals, Fourier sums, and coefficient assembly;
• return a scalar objective such as J = wA*Lambda_A**2 + wB*Lambda_B**2 + wS*Lambda_S**2.

Practical Enzyme API shape

Suggested new area:

src/enzyme/
  rabe_enzyme_api.F90
  rabe_fixed_branch_kernel.F90
  rabe_fixed_quadrature.F90
  rabe_branch_data.F90

test/enzyme/
  enzyme_rabe_driver.c
  test_enzyme_rabe_gradient.py or .sh

The first kernel should be small and explicit, for example:

subroutine rabe_fixed_branch_objective(n, x, branch, weights, J) bind(C)
  use iso_c_binding
  integer(c_int), value :: n
  real(c_double), intent(in)  :: x(n)
  real(c_double), intent(in)  :: branch(*)
  real(c_double), intent(in)  :: weights(*)
  real(c_double), intent(out) :: J

  ! No allocation, no I/O, no error stop.
  ! Compute Lambda_A, Lambda_B, Lambda_S on a fixed branch.
end subroutine

C/C++ wrapper pattern:

extern int enzyme_const;
extern int enzyme_dup;
extern void __enzyme_autodiff(void*, ...);

extern void rabe_fixed_branch_objective(
    int n, const double* x, const double* branch,
    const double* weights, double* J);

void rabe_fixed_branch_gradient(
    int n, const double* x, double* dx,
    const double* branch, const double* weights, double* J)
{
    double dJ = 1.0;
    for (int i = 0; i < n; ++i) dx[i] = 0.0;

    __enzyme_autodiff((void*)rabe_fixed_branch_objective,
        enzyme_const, n,
        enzyme_dup, x, dx,
        enzyme_const, branch,
        enzyme_const, weights,
        enzyme_dup, J, &dJ);
}

Exact calling details should be validated against the selected Enzyme version. The important rule is that branch data and integer indices are inactive, while active arrays are passed with initialized shadow arrays.

Related dependency

libneo work should be tracked separately in itpplasma/libneo#276. That issue covers fixed-cell spline and VMEC/Boozer evaluation kernels that rabe can later call from this Enzyme path.

[krystophny]

Hybrid FD/AD split note

This issue fits a mixed differentiability strategy.

A single solver can expose AD or analytical products for branch-stable subkernels while finite differences or zeroth-order samples remain around hard outputs and driver logic.

Pattern:

• use AD or analytical products for cheap local pieces: profiles, splines, interpolation, field kernels, residual kernels, or smooth proxy maps;
• keep finite differences or zeroth-order sampling for full executable calls, root and topology changes, orbit-class events, convergence logic, restart logic, or long-time loss labels;
• use cheap derivative products to choose reduced finite-difference directions or anisotropic random directions for the expensive target;
• report event changes as diagnostics, not as failed gradients.

Cost model: cheap proxy Jacobians shape where expensive evaluations are spent. This gives useful optimization signal before full end-to-end differentiability exists.

[krystophny]

Branch-margin refinement for the Enzyme / fixed-branch RABE plan:

The fixed-branch AD path should expose explicit branch-margin residuals for the topology choices that are currently treated as inactive branch data. These margins can be used by SensOpt or another optimizer to avoid or diagnose branch changes.

Candidate RABE margins:

m_fieldline_max_gap      = gap between selected two maxima and competing maxima on a field line
m_split_maxima_margin    = smooth distance to the split-maxima condition
m_global_Bmax_gap        = gap between active global B_max line and next-largest B_max
m_minmax_active_gap      = gap between active min/max entries and competitors
m_extremum_curvature     = |d2B/dphi2| at selected maxima/minima
m_eta_b_boundary         = distance of eta_b to problematic endpoint / well boundary
m_quad_endpoint_margin   = minimum transformed quadrature endpoint margin
m_stencil_cell_margin    = distance to VMEC/Boozer interpolation cell or stencil switch

Use in SensOpt form:

r_p = [Lambda_A/B/S proxy residuals, branch margins]

or expose as diagnostics if they should not be protected directly.

Policy:

• Enzyme differentiates the selected branch; it does not differentiate topology changes.
• Branch IDs and branch margins should be returned together with gradients.
• A gradient check should state whether the selected maxima, global Bmax, and interpolation cells stayed fixed under the perturbation.
• If the optimizer should avoid well-topology changes, include the corresponding margin in the protected residual stack.

This should also connect to pyrabe: the JAX port can use the same margin definitions for fixed-branch AD and for deciding when a gradient is branch-local versus globally meaningful.

[krystophny]

Phase-aware SensOpt refinement:

For RABE, the natural region graph is a graph of fixed field-line/well topology and active extrema choices. This complements the fixed-branch Enzyme plan and connects it to itpplasma/sensopt#30.

Suggested RABE region node:

region = (selected two maxima per field line,
          split_maxima flags,
          active global B_max line,
          active min/max entries,
          quadrature interval topology,
          VMEC/Boozer interpolation cell/stencil IDs)

Suggested region edges:

selected maximum changes
new maximum appears / maximum disappears
split_maxima flag changes
active global B_max switches to another field line
eta_b boundary changes
quadrature endpoint/well topology changes
interpolation cell/stencil switches

Each edge has a scalar margin, for example:

m_fieldline_max_gap
m_split_maxima_margin
m_global_Bmax_gap
m_extremum_curvature
m_eta_b_boundary
m_quad_endpoint_margin
m_stencil_cell_margin

Layered use:

1. Inside one region, compute fixed-branch AD or JAX/Enzyme derivatives.
2. Add selected margins to r_p when an optimizer should avoid topology changes.
3. If a topology change may be useful, probe the margin by solving m_i(x + alpha d) = 0, evaluate both sides, and if the branch signature changes, create a new local region and restart fixed-branch optimization there.

This prevents overclaiming global differentiability while still making branch boundaries usable for optimization. RABE/pyrabe can also use the same region signatures so Fortran and JAX paths report compatible branch-locality diagnostics.

[krystophny]

Derivative-source and FD/event-probe refinement for the RABE Enzyme plan:

RABE is a good example where analytical derivatives, JAX/Enzyme AD, finite differences, and event diagnostics should be combined rather than forcing one derivative mode through the whole code.

Analytical derivatives first

Use hand/analytical derivatives for structured smooth operations:

Fourier / Boozer basis sums
linear dependence on Boozer coefficients
fixed-cell field evaluation
fixed Gauss-Legendre quadrature sums
surface / field-line sums and weighted averages
simple algebra in Lambda_A/B/S assembly

This keeps branch assumptions explicit and avoids unnecessary AD overhead.

Reverse AD / Enzyme target

Reverse mode is appropriate for scalar outputs with many active coefficients, e.g.

J = wA Lambda_A^2 + wB Lambda_B^2 + wS Lambda_S^2

or one scalar Lambda_A(s) with respect to many Boozer coefficients, but only after fixing the branch topology:

same selected maxima per field line
same active global Bmax
same split_maxima status
same quadrature/well topology
same interpolation cells/stencils

Forward-mode / directional products

Forward-mode or tangent-linear products are useful along a small number of SensOpt or UQ directions:

J_margin * d
J_Lambda * d

especially for branch-margin diagnostics and finite-difference validation.

FD and event-probe policy

Branch-preserving FD:

same branch signature on p +/- eps d

can validate JAX/Enzyme/analytic derivatives.

Cross-boundary FD should be classified before use:

selected maximum switch -> kink or event
split_maxima change -> topology event
active global Bmax switch -> active-branch kink
well creation/annihilation -> topology transition

Margin-guided boundary probes should solve

m_i(p + alpha d) = 0

for margins such as:

m_fieldline_max_gap
m_split_maxima_margin
m_global_Bmax_gap
m_extremum_curvature
m_eta_b_boundary
m_quad_endpoint_margin
m_stencil_cell_margin

then evaluate both sides and classify smooth/kink/jump/event.

Relation to pyrabe

The same policy should be mirrored in pyrabe:

jax.grad = fixed-branch derivative
branch_margins = validity diagnostics
cross-boundary FD = event classification / region graph edge

This gives a clean route from JAX fixed-branch gradients to Fortran/Enzyme gradients without overclaiming global differentiability.

[GeorgGrassler]

Currently investigating analytical adjoint approach as only appropriate for the fourier_field_t case and enzym struggels with polymorphism in the main tool chain.

[krystophny]

Automated triage pass (Claude), grounded in the repository at 794bdfd. Verify before acting.

Still relevant, but the issue text has drifted from the tree in three ways and the six comments have grown it into a tracker. Below is what actually changed, followed by a spec for the one milestone that is implementable today.

Drift 1 — obstacle 6 is void. src/vmec/boozer_converter.F90 was deleted in e1142da ("Use libneo boozer_sub; drop rabe local converter (#100)", 893 lines removed); src/vmec/ now contains only boozer_field.f90 and CMakeLists.txt. The "differentiable VMEC-to-Boozer step through the existing boozer_converter Fortran source" no longer has a source in this repo. That workstream is entirely itpplasma/libneo#276.

Drift 2 — obstacle 1 misdescribes the code. src/utils/find_extrema.f90 has no Newton refinement to apply the implicit-function theorem to. refine_maximum (lines 96–134) iterates an 11-point linspace + maxloc subdivision (n_refine = 11, maxloc at line 112) and returns a grid node, loc = x_sub(i_max) (line 133). nudge_maxima_inward (src/fieldline/make_fieldline.f90:381-403) then shifts that node by the bisection's own error estimate (lines 385-386, phi_max(1) + phi_max_error(1) / phi_max(2) - phi_max_error(2)). So d phi_max/dp is not merely "the derivative of the wrong branch" — it is the derivative of a bracket-termination artefact. IFT refinement has to be written, not differentiated around.

Drift 3 — the build prerequisite is partly done, on a branch nobody linked here. origin/build/llvm_light (@GeorgGrassler, last commit 1efb21c, 2026-07-14) already carries: an LLVMFlang arm in cmake/compiler.cmake (lines 8 and 32), .github/workflows/flang-smoke-test.yml (nightly cron: '0 3 * * *', installs LLVM 19 and runs with FC=flang-new-19), cmake/netcdf.cmake building netcdf-fortran from source, and — most importantly — src/offset/ad_driver.f90, a single-entry-point rabe_fourier_offset_coefficients(m, n, B_mn, ..., lambda_a, lambda_b, nu_star_crit) whose docstring names Enzyme as the intended consumer ("a single, well-defined target for the eventual Enzyme adjoint"). That file is the target; it just is not fixed-branch yet (it calls fourier_field_init and make_flock_of_fieldlines, i.e. maxima discovery, inside the would-be active region).

The branch is 1 commit behind and 5 ahead of main (git rev-list --left-right --count main...origin/build/llvm_light → 1 5); its merge base is 5f035b7, so 4f70aa4 and 5f035b7 are both already in it. Only 794bdfd ("update: paper number") is missing — this is a trivial fast-forward, not a stale fork.

---

Purpose

Stellarator shape optimization needs d Lambda_{A,B,S} / d B_mn at a cost independent of the number of DOFs, which is the reverse-mode cost model. Today the only gradient path is pyrabe's JAX reimplementation, which must be kept numerically in sync with the Fortran kernel by hand, and README.md:254-259 already advertises "## Adjoint Gradients — Status: under active development." against no source in this repo (rg -i 'enzyme|autodiff' src/ returns nothing). A validated Enzyme gradient on one scalar, one surface, one fixed well topology removes the duplication risk and gives SensOpt a usable descent direction; everything else in this issue is downstream of that one number being right.

Current state

origin/build/llvm_light:src/offset/ad_driver.f90 is the entry point:

call fourier_field_init(field, m, n, B_mn, B_theta_covariant, B_phi_covariant, &
                        nfp, n_grid)
call make_flock_of_fieldlines(flock, max_n_fieldlines, iota, field, &
                              M_pol, N_tor, real(nfp, dp))
call calc_offset_coefficients(flock, R, dr_dAtheta, lambda_a, lambda_b)
nu_star_crit = calc_nu_star_crit(flock, R)

Everything Enzyme cannot handle sits between lines 2 and 3 of that snippet. Concretely, on main:

• Dynamic dispatch. Every field evaluation goes through the abstract field_t deferred bindings (src/field/field_base.f90:12-17), called as class(field_t) in src/fieldline/fieldline_integrands.f90:16,33,44,70-71. This is the "enzym struggels with polymorphism" @GeorgGrassler reported on 2026-07-06.
• Module-global mutable state and procedure pointers on the entire integration path: src/utils/integrate.f90:18 and :165-166; src/utils/find_extrema.f90:13; src/fieldline/fieldline_integrals.f90:8-9; src/fieldline/field_along_fieldline.f90:9-10. The developer note at integrate.f90:99-112 states this is deliberate and that "The call chain can not be made thread safe." The same property defeats Enzyme's activity analysis through the indirect call.
• Adaptive quadrature. integrate_1d calls QUADPACK dqag (aliased qag, integrate.f90:4) in quadkind = 8 (:8) with rel_error_tol_quadkind = 1.0e-6 (:28), the call itself at :54-68, through the fetched third-party sources (cmake/quadpack.cmake). The substitution wrapper integrate_1d_substituted (:172-207) already applies t**2 = x - a / t**2 = b - x (documented at :185-193), so only the rule has to be swapped.
• Hard selections. get_global_B_max (src/fieldline/make_fieldline.f90:405-415, two maxval calls at :412-413), minval(flock%fieldlines%B_max(1)) in calc_nu_star_crit (src/offset/coefficients.f90:173), and pick_maximum's maxloc/minloc (make_fieldline.f90:373,378), reached from pick_maximum_on_each_side (:335-354).
• 95 error stop sites in src/ (rg -c 'error stop' src/ | awk -F: '{s+=$2} END {print s}' → 95), several directly on the milestone-1 path: find_extrema.f90:85,89, make_fieldline.f90:143,201, integrate.f90:45,94, and fieldline_integrands.f90:59 — note this last one is an abort, not a clamp: calc_lambda_squared prints and error stops when 1 - B*eta < 0. pyrabe's max(1 - B*eta_b, 0) clamp has no counterpart in the Fortran.

One property worth exploiting: fourier_field_t is exactly linear in B_mn. src/field/fourier_field.f90:104-105 builds grid_B as sum(B_mn*cos(m*theta - n*nfp*phi)), and construct_splines_2d (libneo src/interpolate/interpolate.f90:276-326 at the pinned SHA 54e372f2, see cmake/libneo.cmake:29) passes the samples to spl_per, whose branching is on spline order only (libneo src/spl_three_to_five.f90:639-688), never on the data. So the interpolant is a linear map of the sample values and

dB(theta,phi)/dB_mn(k) = spline_of( cos(m_k*theta - n_k*nfp*phi) ) evaluated at (theta,phi)

is available in closed form. That gives an analytic reference for the innermost derivative, independent of both Enzyme and finite differences.

Reproduction / verification recipe

Fixture: the 5-mode anti-sigma field in python/example_fourier.py. Constants at lines 113-122 (EPS_1 = 0.005, EPS_2 = 0.00009375, IOTA = 0.47, NFP = 10, B_0 = 1.0, B_PHI_COV = 8.0, B_THETA_COV = 0.0, R = 8.0, PSI_TOR_EDGE = 2.812e-07); modes at lines 124-128 (m = [0,1,1,2,3], n = [0,0,1,1,1], B_mn = B_0*[1.0, -EPS_2, -EPS_2*0.5, -EPS_1, -EPS_2*0.5]). Note lines 129-130 then perturb the fixture: delta_B_1 = B_0*EPS_2*0.01; B_mn[1] += delta_B_1. Reuse the file's own post-perturbation B_mn verbatim rather than reconstructing it from the constants.

calc_lambda_a_analytic (lines 21-29) gives a small-aspect-ratio asymptotic estimate of Lambda_A (its own docstring says so), not an exact closed form. It is a sanity bound on the primal, not a tolerance-grade reference. The tolerance-grade primal reference is the current gfortran output of the example itself.

Baseline today (gfortran, main):

cd /path/to/rabe
make python           # cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_PYTHON_BINDINGS=ON
python python/example_fourier.py     # prints Lambda_A, Lambda_B, nu_star_crit

There is no gradient to compare against; that is the gap. Central differences on B_mn(4) (Fortran 1-based; m=2, n=1, the dominant EPS_1 mode) with h = 1e-6*B_0 are the interim reference.

Expected end state. Do not try to pass the option through make: Makefile:5-10 forwards only LIBNEO_REF and LIBNEO_PATH from the command line, and build/CMakeCache.txt is a file prerequisite, so a stale cache is never reconfigured. Configure explicitly:

rm -rf build
FC=flang-new-19 cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DRABE_ENABLE_ENZYME=ON
cmake --build build
ctest --test-dir build/test -R TestEnzymeLambdaAGradient --output-on-failure

with the printed max relative deviation between dLambda_A/dB_mn from Enzyme and from central differences below 1e-5 (loosened from the 1e-6 quadrature tolerance because FD itself carries sqrt(eps)/h noise).

Behaviour (BDD)

Scenario: Analytic field-linearity check on the innermost derivative
  Given the 5-mode anti-sigma coefficients read from python/example_fourier.py:124-130
        (m = [0,1,1,2,3], n = [0,0,1,1,1], B_mn after the delta_B_1 perturbation)
  And n_grid = 200, nfp = 10
  When test/unit/test_fourier_field_linearity.x evaluates
       B(theta=0.3, phi=0.07) for B_mn and for B_mn + h*e_k, h = 1e-8
  Then (B(B_mn + h*e_k) - B(B_mn))/h equals the spline of cos(m_k*theta - n_k*nfp*phi)
       at the same point to a relative tolerance of 1e-9 for every k in 1..5

Scenario: Fixed Gauss-Legendre well integrals reproduce the adaptive result
  Given a fieldline%phi_max pair produced by make_flock_of_fieldlines on the
        anti-sigma fixture, dumped from the first field line of the flock
  When integrate_1d_substituted_gl(wrapper_lambda_over_B_squared, a, b, res, n_nodes=64)
       is compared against integrate_1d_substituted (src/utils/integrate.f90:172)
       on the same integrand and the same [a, b]
  Then the two results agree to relative 1e-7, i.e. below the dqag setting
       rel_error_tol_quadkind = 1.0e-6 at src/utils/integrate.f90:28

Scenario: Enzyme reverse gradient of Lambda_A matches central differences
  Given the fixed-branch kernel rabe_fixed_branch_lambda_a with a branch record
        captured from a prior discovery run on the same B_mn
  When test/enzyme/enzyme_rabe_driver.c calls __enzyme_autodiff with enzyme_dup on B_mn
  Then dLambda_A/dB_mn(k) agrees with central differences (h = 1e-6*B_0) to
       relative 1e-5 for all k, and no NaN appears in the shadow array

Scenario: Branch change is refused, not silently differentiated
  Given a branch record whose selected phi_max indices were captured at B_mn
  When the kernel is called at B_mn + delta where the coarse argmax index changes
  Then rabe_fixed_branch_lambda_a returns branch_ok = 0 and leaves Lambda_A unchanged
  And the driver reports the offending field line index rather than a gradient

Implementation plan

Ordering: steps 1–2 are prerequisites already half-done on origin/build/llvm_light; land that branch first.

1. Merge origin/build/llvm_light build wiring into main — cmake/compiler.cmake (LLVMFlang arms at lines 8 and 32), cmake/netcdf.cmake, .github/workflows/flang-smoke-test.yml. The branch is only one commit behind main, so git rebase main (or a plain merge) is a fast-forward of 794bdfd onto it; no conflict resolution against 4f70aa4/5f035b7 is needed, they are already in the branch.

2. CMakeLists.txt — add option(RABE_ENABLE_ENZYME "Build the Enzyme AD path" OFF); when ON and CMAKE_Fortran_COMPILER_ID MATCHES "LLVMFlang", append -fpass-plugin=${ENZYME_PLUGIN} to the new rabe_enzyme_lib target only, with find_file for LLVMEnzyme-*.so and a message(FATAL_ERROR) if unset. Do not touch global flags — the golden record must keep building with the gfortran 11 pin in flake.nix.

3. Create src/enzyme/rabe_branch_data.F90 — derived type branch_record_t holding, per field line: phi_max_index(2) (coarse-grid indices), line_of_global_B_max, line_of_min_B_max (the minval winner from src/offset/coefficients.f90:173), split_maxima, and n_fieldlines. Plus capture_branch(flock, record) and logical function branch_matches(flock, record).

4. Create src/enzyme/rabe_fixed_quadrature.F90 — integrate_1d_substituted_gl(f, a, b, n_nodes, result) with hard-coded Gauss–Legendre nodes/weights (n_nodes = 64), reusing the t**2 = x - a / t**2 = b - x split documented at src/utils/integrate.f90:185-193. No procedure pointer: take the integrand as an explicit procedure() dummy, or better, inline the four integrands as separate specialised routines so the call is direct. The four integrands to port are the contained wrappers in src/fieldline/fieldline_integrals.f90:48-102 (wrapper_lambda_over_B_squared, wrapper_local_radial_drift, wrapper_one_over_B_squared, wrapper_nabla_s_over_B_squared), which delegate to src/fieldline/fieldline_integrands.f90.

5. Create src/enzyme/rabe_fixed_branch_kernel.F90 — subroutine rabe_fixed_branch_lambda_a(mn_max, m, n, B_mn, spline_coeff, branch, R, dr_dAtheta, Lambda_A, branch_ok). Rules: type(fourier_field_t) concrete, never class(field_t); no save; no allocation; no error stop (replace with branch_ok = 0 and an IEEE quiet NaN); integers and branch inactive. Reimplement the well integrals against rabe_fixed_quadrature, and the surface average as the plain sum(...)*dxi_0/normalization expressions of src/fieldline/surface_average.f90:43-57.

6. Add IFT refinement in src/enzyme/rabe_fixed_branch_kernel.F90 — given branch%phi_max_index, run a fixed count (3) of Newton steps on dB/dphi = 0 using the compute_B_and_dB_dx binding (src/field/field_base.f90:13) plus a finite-difference second derivative, and replace nudge_maxima_inward entirely on this path. Do not reuse find_extrema.f90.

7. Create src/enzyme/rabe_enzyme_api.F90 — bind(C) shim with iso_c_binding scalars only (integer(c_int), value), forwarding to step 5. This is the __enzyme_autodiff target.

8. Create src/enzyme/CMakeLists.txt and add add_subdirectory(enzyme) to src/CMakeLists.txt (after add_subdirectory(field)), guarded by RABE_ENABLE_ENZYME. Link field_lib and interpolate. There is no fourier_field_lib — fourier_field.f90 is compiled into field_lib (src/field/CMakeLists.txt:1), and field_lib links interpolate PRIVATE, so any new target that touches SplineData2D must name interpolate itself.

9. Create test/enzyme/enzyme_rabe_driver.c — the wrapper from the 2026-05-23T11:30 comment, with enzyme_const on mn_max, m, n, branch, R, dr_dAtheta and enzyme_dup on B_mn/dB_mn and on Lambda_A/dLambda_A. Zero dB_mn before the call.

10. Create test/enzyme/test_enzyme_lambda_a_gradient.py — loads the fixture of python/example_fourier.py:113-130, calls the C driver via ctypes, computes central differences through the same C entry point, asserts relative agreement < 1e-5. Register in test/enzyme/CMakeLists.txt as add_test(TestEnzymeLambdaAGradient ...) with set_tests_properties(TestEnzymeLambdaAGradient PROPERTIES LABELS "slow") — it must not enter make test, which is ctest -L quick (Makefile:24-25) and runs under gfortran. Add add_subdirectory(enzyme) to test/CMakeLists.txt under the same guard.

11. Create test/unit/test_fourier_field_linearity.f90 and register it in test/unit/CMakeLists.txt following the existing four-line idiom, e.g.

    add_executable(test_fourier_field_linearity.x test_fourier_field_linearity.f90)
    target_link_libraries(test_fourier_field_linearity.x field_lib utils_lib interpolate)
    add_test(TestFourierFieldLinearity test_fourier_field_linearity.x)
    set_tests_properties(TestFourierFieldLinearity PROPERTIES LABELS "quick")

    This one is compiler-independent and guards step 5's premise even without flang.

12. Update README.md:254-259 to state what actually exists once step 10 passes, and CLAUDE.md's library tree (lines 32-60), which is already stale: offset_lib (src/offset/) is missing entirely, and field_lib is listed with field_base.f90 only, omitting fourier_field.f90.

Sibling repo: nothing in libneo is required for milestone 1 — fourier_field_t only needs construct_splines_2d/evaluate_splines_2d_der (src/field/fourier_field.f90:5,109,145), which are pure arithmetic once built. itpplasma/libneo#276 (fixed-cell spline evaluators) becomes a hard dependency only at milestone 2, when the input moves to boozer_field_t and libneo's boozer_sub.

Acceptance criteria

• FC=flang-new-19 cmake -S . -B build -DCMAKE_BUILD_TYPE=Release && cmake --build build succeeds on main, and the nightly flang-smoke-test.yml (which runs plain make + make test, i.e. Debug and -L quick) is green.
• ctest --test-dir build/test -R TestFourierFieldLinearity passes under gfortran, relative tolerance 1e-9.
• ctest --test-dir build/test -R TestEnzymeLambdaAGradient passes with -DRABE_ENABLE_ENZYME=ON; max relative Enzyme-vs-central-difference deviation over the 5 modes < 1e-5, printed by the test.
• rabe_fixed_branch_lambda_a reproduces the gfortran python/example_fourier.py Lambda_A to relative 1e-7 — the fixed 64-node Gauss–Legendre rule must not shift the primal beyond the 1e-6 dqag setting.
• rg -c 'error stop' src/enzyme/ returns nothing; no save, no allocate, no class( in src/enzyme/.
• make test_all and make golden still pass under the pinned gfortran-11 flake.nix toolchain. Note the golden gate is rtol=1e-10, atol=0.0 via np.testing.assert_allclose (test/golden/compare_golden.py:34-39), not bit-exactness — no test/golden/expected/rabe.nc regeneration should be needed.

References

• Moses & Churavy, Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients, NeurIPS 2020 — the Enzyme paper cited in the issue.
• Moses, Churavy, Paehler, Hückelheim, Narayanan, Schanen, Doerfert, Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme, SC'21 — covers the LLVM-IR activity analysis that the module-global procedure pointers above defeat.
• Hascoët & Pascual, The Tapenade Automatic Differentiation Tool: Principles, Model, and Specification, ACM TOMS 39(3), 2013 — the source-transform alternative for Fortran, worth a paragraph in the design note if flang+Enzyme stalls.
• Piessens, de Doncker-Kapenga, Überhuber & Kahaner, QUADPACK: A Subroutine Package for Automatic Integration, Springer 1983 — dqag's adaptive subdivision, the reason step 4 replaces it rather than differentiating it.
• Antonsen, Paul & Landreman, Adjoint approach to calculating shape gradients for three-dimensional magnetic confinement equilibria, J. Plasma Phys. 85, 905850207 (2019) — the analytic-adjoint route @GeorgGrassler raised on 2026-07-06.
• Dudt & Kolemen, DESC: A stellarator equilibrium solver, Phys. Plasmas 27, 102513 (2020); Landreman et al., SIMSOPT: A flexible framework for stellarator optimization, JOSS 6(65), 3525 (2021) — the two codes solving the same gradient problem. DESC's bounce-integral objectives cap the number of detected wells with a fixed num_well argument, which is the same fixed-branch idea as step 3; the exact correspondence to steps 3 and 6 was not re-verified here and should be checked against DESC's current desc.integrals API before it is relied on.
• Helander, Geiger & Maaßberg, On the bootstrap current in stellarators and tokamaks, Phys. Plasmas 18, 092505 (2011); Landreman & Catto, Omnigenity as generalized quasisymmetry, Phys. Plasmas 19, 056103 (2012) — the physics the differentiated quantities encode.

The fourier_field_t linearity argument and the branch-record design in steps 3–6 have no specific literature reference; they are implementation details internal to this code.

Scope and risk

Milestone 1 alone is 2–3 PRs (build wiring; fixed quadrature + linearity test; kernel + Enzyme driver), not one. The main risk is that flang-new 19 plus Enzyme cannot handle the type(SplineData2D) allocatable component at all, in which case step 5 must additionally flatten spl%coeff into a plain real(dp) array argument — decide that with a one-day spike before writing the kernel. Nothing here touches the gfortran golden path, so a failure is contained.

Suggested split. This issue should be reduced to the milestone-1 spec above and the rest spun out: (a) branch-margin residuals m_fieldline_max_gap, m_global_Bmax_gap, m_split_maxima_margin etc. (2026-05-23T23:24 comment) belong in a rabe issue of their own, cross-linked to itpplasma/sensopt#30; (b) the hybrid FD/AD and region-graph policy comments (2026-05-23T22:54, 23:46, 2026-05-24T04:42) are org-level design and belong in sensopt, not here; (c) milestone 2 (input boundary at boozer_field_t/libneo boozer_sub) is blocked on itpplasma/libneo#276 and should be a separate issue rather than a bullet. Related, and complementary rather than duplicate: itpplasma/NEO#1 (same fixed-branch pattern for NEO — should share the branch-record vocabulary) and itpplasma/pyrabe#23 (the JAX-side counterpart). Neither cross-repo issue was inspected in this pass; confirm their scope before linking them as dependencies.