Equations and derivations

This page states the explicit equations solved (or approximated) by vmec_jax and connects them to the VMEC2000 formulation. The goal is to make the physics, the coordinate conventions, and the force-balance residuals fully transparent so that parity checks can be done equation-by-equation.

Notation and conventions

We use VMEC’s curvilinear flux coordinates:

  • \(s \in [0,1]\): normalized toroidal flux label (VMEC radial coordinate).

  • \(\theta \in [0,2\pi)\): poloidal angle.

  • \(\zeta \in [0,2\pi)\): field-period toroidal angle.

  • \(\phi_{\mathrm{phys}} = \zeta/\mathrm{NFP}\): physical toroidal angle.

The Fourier phase convention is:

\[\mathrm{phase}(m,n;\theta,\zeta) = m\theta - n\zeta,\]

where \(n\) is the field-period toroidal mode number (VMEC stores \(xn = n\,\mathrm{NFP}\) in wout).

VMEC weighted DFT tables (fixaray)

VMEC does not use a plain FFT for its force/residual transforms. Instead, fixaray builds weighted trig tables on a symmetry-aware grid and applies explicit normalization factors. Let \(\theta_i\) be the VMEC theta grid over \([0,\pi]\) (with endpoint half-weights) and \(\zeta_k\) the zeta grid over one field period. VMEC defines

\[\mathrm{dnorm} = \frac{1}{n_\zeta\,(n_{\theta2}-1)},\qquad \mathrm{mscale}_0 = 1,\quad \mathrm{mscale}_{m>0}=\sqrt{2},\]

and the weighted cosine table

\[\mathrm{cosmui}_{i,m} = \mathrm{dnorm}\,w_i\,\mathrm{mscale}_m \cos(m\theta_i),\]

with \(w_0=w_{n_{\theta2}-1}=1/2\) and \(w_i=1\) elsewhere. The sine table is defined analogously, with the same weights and mscale. Zeta tables use nscale (also \(\sqrt{2}\) for \(n>0\)) and, for derivative terms, include the field-period multiplier \(n\,\mathrm{NFP}\):

\[\mathrm{cosnvn}_{k,n} = (n\,\mathrm{NFP})\,\mathrm{cosnv}_{k,n}, \qquad \mathrm{sinnvn}_{k,n} = -(n\,\mathrm{NFP})\,\mathrm{sinnv}_{k,n}.\]

vmec_jax uses these tables in tomnsps so that the Fourier-space force arrays exactly match VMEC2000. See References [4-6] for the original VMEC2000 tables and the VMEC++ DFT/basis discussion.

Two-stage DFT for tomnsps

VMEC’s tomnsps uses a separable real basis in \(\theta\) and \(\zeta\). For a real-space kernel \(F(\theta_i,\zeta_k)\) defined on the VMEC grid, the weighted theta projection is

\[\tilde F^{(c)}_{m}(\zeta_k) = \sum_{i=0}^{n_{\theta2}-1} F(\theta_i,\zeta_k)\,\mathrm{cosmui}_{i,m},\]
\[\tilde F^{(s)}_{m}(\zeta_k) = \sum_{i=0}^{n_{\theta2}-1} F(\theta_i,\zeta_k)\,\mathrm{sinmui}_{i,m}.\]

The zeta projection then yields the Fourier coefficients

\[F^{cc}_{m,n} = \sum_{k=0}^{n_\zeta-1} \tilde F^{(c)}_{m}(\zeta_k)\,\mathrm{cosnv}_{k,n},\]
\[F^{ss}_{m,n} = \sum_{k=0}^{n_\zeta-1} \tilde F^{(s)}_{m}(\zeta_k)\,\mathrm{sinnv}_{k,n}.\]

Derivative terms in VMEC use the scaled tables \(\mathrm{cosnvn}_{k,n}=(n\,\mathrm{NFP})\,\mathrm{cosnv}_{k,n}\) and \(\mathrm{sinnvn}_{k,n}=-(n\,\mathrm{NFP})\,\mathrm{sinnv}_{k,n}\). In vmec_jax we therefore compute the same base transforms and apply the analytic factor \(n\,\mathrm{NFP}\) after the zeta contraction for the derivative blocks. This reduces the number of dot-product contractions while preserving VMEC2000 parity exactly.

Implementation detail: the theta contractions for multiple force kernels are stacked into a single batched dot_general call (GEMM), and the zeta contractions are likewise stacked by basis type (cosine vs sine). This follows the separable product identities (see Eqs. 5.55–5.56 in the VMEC++ numerics notes) while keeping the VMEC2000 normalization and parity masks intact.

Ideal MHD equilibrium

The ideal MHD equilibrium is defined by:

\[\nabla p = \mathbf{J} \times \mathbf{B},\]

with Maxwell’s equations (in magnetostatic form):

\[\nabla \cdot \mathbf{B} = 0, \qquad \mathbf{J} = \frac{1}{\mu_0}\nabla \times \mathbf{B}.\]

The pressure is a flux function: \(p = p(s)\) and is specified by the VMEC input profiles. VMEC (and vmec_jax) use pressure in units of \(\mu_0\,\mathrm{Pa}\) so that \(p\) has the same units as \(B^2\).

Energy principle (VMEC formulation)

VMEC solves for a stationary point of the ideal-MHD energy functional in straight-field-line coordinates:

\[W = \frac{1}{(2\pi)^2}\int \left(\frac{B^2}{2} + \frac{p}{\gamma-1}\right) dV,\]

where \(\gamma\) is the ratio of specific heats (VMEC input GAMMA). The VMEC fixed-boundary update loop can be viewed as a (preconditioned) steepest-descent method that drives the force residuals to zero. See References [1-3] for the original VMEC formulation.

Flux coordinates and straight-field-line angle

VMEC introduces a scalar field \(\lambda(s,\theta,\zeta)\) to define the straight-field-line poloidal angle:

\[u = \theta + \lambda(s,\theta,\zeta).\]

Field lines are straight in \((u,\zeta)\):

\[\frac{du}{d\zeta} = \iota(s),\]

where \(\iota(s)\) is the rotational transform.

Internal scaling and regularity (scalxc)

VMEC enforces regularity at the magnetic axis by storing odd-m contributions in an internal form that factors out \(\sqrt{s}\):

\[X(s,\theta,\zeta) = X_{\mathrm{even}}(s,\theta,\zeta) + \sqrt{s}\,X_{\mathrm{odd,int}}(s,\theta,\zeta).\]

Equivalently,

\[X_{\mathrm{odd,int}} = \frac{X_{\mathrm{odd,phys}}}{\sqrt{s}}.\]

VMEC implements this via the scalxc array, which is 1 for even-m harmonics and \(1/\sqrt{s}\) for odd-m harmonics. scalxc is applied when interpolating coefficients between radial grids and when assembling preconditioned residuals (VMEC2000 profil3d / interp / scalxc).

On the axis, VMEC applies odd-m rules:

  • \(m=1\): extrapolate the internal odd field to the axis by copying the first off-axis value,

  • \(m\ge 2\): force the internal odd field to zero on-axis.

m=1 internal constraint (lconm1)

When LCONM1 is enabled (VMEC default for 3D runs), VMEC stores the m=1 boundary coefficients in a constrained internal basis:

\[R^{s}_{1n,\mathrm{int}} = \frac{1}{2}\left(R^{s}_{1n,\mathrm{phys}} + Z^{c}_{1n,\mathrm{phys}}\right), \qquad Z^{c}_{1n,\mathrm{int}} = \frac{1}{2}\left(R^{s}_{1n,\mathrm{phys}} - Z^{c}_{1n,\mathrm{phys}}\right).\]

This transformation is applied in VMEC2000 readin and inverted when converting to physical coefficients for diagnostics. vmec_jax uses the same internal basis so that boundary handling and multigrid interpolation match VMEC2000.

Magnetic field representation

In VMEC’s flux-coordinate representation, the magnetic field has no radial contravariant component:

\[B^s = 0.\]

VMEC therefore stores only the contravariant components in the angular directions:

\[\mathbf{B} = B^u \nabla u + B^v \nabla v, \qquad v \equiv \zeta.\]

In terms of VMEC’s flux functions \(\Phi(s)\) (toroidal flux) and \(\chi(s)\) (poloidal flux), we define:

\[\Phi'(s) \equiv \frac{d\Phi}{ds} \quad (\text{``phipf''}), \qquad \chi'(s) \equiv \frac{d\chi}{ds} \quad (\text{``chipf''}).\]

VMEC’s contravariant components (bsupu and bsupv in wout) are computed as:

\[B^v = \frac{\Phi'(s) + \mathrm{lamscale}\,\partial_{\theta}\lambda} {\mathrm{signgs}\,\sqrt{g}\,2\pi},\]
\[B^u = \frac{\chi'(s) - \mathrm{lamscale}\,\partial_{\zeta}\lambda} {\mathrm{signgs}\,\sqrt{g}\,2\pi}.\]

Here:

  • \(\sqrt{g}\) is the signed Jacobian,

  • signgs is VMEC’s sign convention such that signgs*sqrtg is positive away from the magnetic axis,

  • lamscale is the VMEC scaling applied to \(\lambda\) derivatives (stored in wout and used by vmec_jax for parity).

The covariant components are defined by:

\[B_i = \mathbf{B}\cdot \mathbf{e}_i = \sum_{j} g_{ij} B^j, \qquad i,j \in \{s,u,v\},\]

where \(g_{ij}\) is the covariant metric and \(\mathbf{e}_i = \partial_i \mathbf{r}\). VMEC stores these as bsub* in wout.

bcovar + add_fluxes (poloidal flux correction)

VMEC updates the contravariant \(B^u\) using the full-mesh poloidal flux function \(\chi'(s)\) (chips). In VMEC2000 add_fluxes, chips is computed from force balance on each surface:

\[\chi'(s) = \frac{I_\varphi(s) - \langle g^{uu} B_u + g^{uv} B_v \rangle} {\langle g^{uu}/\sqrt{g} \rangle},\]

where the angle brackets denote the VMEC surface quadrature, and \(I_\varphi(s)\) is the integrated toroidal current (icurv). VMEC then applies the correction

\[B^u \leftarrow B^u + \chi'(s)\,\frac{1}{\sqrt{g}}.\]

VMEC stores the half-mesh averaged chipf in wout; vmec_jax follows VMEC’s averaging rules to convert between chipf and chips.

Current density

The current density follows directly from the curl:

\[\mathbf{J} = \frac{1}{\mu_0} \nabla \times \mathbf{B}.\]

VMEC reports covariant current components in wout as jcuru and jcurv (poloidal and toroidal current densities on the half mesh) and uses these in the force kernels. The parallel and perpendicular currents satisfy:

\[\mathbf{J} = \mathbf{J}_{\parallel} + \mathbf{J}_{\perp}, \qquad \nabla p = \mathbf{J}_{\perp} \times \mathbf{B}.\]

For optimization diagnostics, vmec_jax also exposes the JXBFORCE real-space current channels as

\[J^\theta = \frac{\texttt{itheta}}{\sqrt{g}}, \qquad J^\zeta = \frac{\texttt{izeta}}{\sqrt{g}},\]

on the full radial mesh. The vj.JVector objective returns these flux-coordinate components flattened over the selected surfaces and angular grid. vj.BVector returns the corresponding Cartesian magnetic-field vector (B_x,B_y,B_z) on one selected radial surface.

Redl Bootstrap-Current Mismatch

For finite-beta stage-one studies, vmec_jax exposes a differentiable Redl bootstrap-current residual. The residual follows the normalized SIMSOPT form

\[R_j = \frac{\langle\mathbf{J}\cdot\mathbf{B}\rangle_{\mathrm{VMEC}}(s_j) - \langle\mathbf{J}\cdot\mathbf{B}\rangle_{\mathrm{Redl}}(s_j)} {\left[\sum_k \left(\langle\mathbf{J}\cdot\mathbf{B}\rangle_{\mathrm{VMEC}}(s_k) + \langle\mathbf{J}\cdot\mathbf{B}\rangle_{\mathrm{Redl}}(s_k) \right)^2\right]^{1/2}}.\]

The Redl term uses polynomial density and temperature profiles in the same ascending-coefficient convention as SIMSOPT ProfilePolynomial. For the standard finite-beta stage-one examples, vj.standard_finite_beta_profiles constructs

\[n_e(s) = n_{e0}(1 - 0.99 s^5), \qquad T_e(s) = T_{e0}(1 - 0.99 s),\]

with ni=ne, Ti=Te, Zeff=1, and \(p(s)=e(n_eT_e+n_iT_i)\) in Pascals. The amplitudes use the same scaling as the SIMSOPT finite-beta/bootstrap examples,

\[n_{e0} = 3\times 10^{20} \left(\frac{\beta/100}{0.05}\right)^{1/3}, \qquad T_{e0} = 15\,\mathrm{keV} \left(\frac{\beta/100}{0.05}\right)^{2/3}.\]

vj.with_pressure_profile converts this pressure profile to VMEC AM and PRES_SCALE input fields while vj.RedlBootstrapMismatch receives the same density/temperature coefficients. The effective trapped-particle fraction is evaluated with fixed Gauss-Legendre quadrature using the substitution \(y = \sqrt{1-\lambda B_{\max}}\), which removes the endpoint singularity in the standard integral

\[f_t = 1 - \frac{3}{4}\langle B^2\rangle \int_0^{1/B_{\max}} \frac{\lambda\,d\lambda} {\left\langle\sqrt{1-\lambda B}\right\rangle}.\]

This differs from SIMSOPT’s post-processing routine, which refines angular extrema with splines. The vmec_jax form is intentionally fixed-shape and differentiable for use inside exact-Jacobian optimization.

Force balance in VMEC (residual form)

VMEC evaluates the force balance in real space, then transforms the residual forces back to Fourier space. In VMEC2000, these residuals are packaged as tomnsps Fourier arrays and projected into three directions:

  • \(F_R\): radial (R) force balance residual,

  • \(F_Z\): vertical (Z) force balance residual,

  • \(F_\lambda\): stream-function (lambda) residual.

These are combined into scalar norms that appear in the VMEC screen output:

\[\mathrm{FSQR} = r_1\,\mathrm{fnorm}\,\lVert F_R \rVert^2, \qquad \mathrm{FSQZ} = r_1\,\mathrm{fnorm}\,\lVert F_Z \rVert^2,\]
\[\mathrm{FSQL} = \mathrm{fnormL}\,\lVert F_\lambda \rVert^2,\]

where fnorm, fnormL, and r1 are the VMEC normalization factors computed from bcovar (half-mesh metrics + bsup/bsub). vmec_jax reproduces these scalars from the same internal quantities to match VMEC2000’s per-iteration printout.

Time-step control (Garabedian update)

VMEC’s nonlinear fixed-boundary iteration uses a Garabedian-style conjugate gradient update with a time-step control mechanism. The update computes

\[\tau_n = \min\left(\left|\ln\frac{\mathrm{fsq}_n}{\mathrm{fsq}_{n-1}}\right|,\; 0.15\right),\]

and maintains a moving average \(\overline{\tau}\) over ndamp steps. The damping factor is

\[\Delta\tau = \frac{\Delta t\,\overline{\tau}}{2}, \qquad b_1 = 1-\Delta\tau, \qquad \mathrm{fac} = \frac{1}{1+\Delta\tau}.\]

The update in VMEC2000 is:

\[\dot{x} \leftarrow \mathrm{fac}\,\bigl(b_1\dot{x} + \Delta t\,F\bigr), \qquad x \leftarrow x + \Delta t\,\dot{x},\]

where \(F\) is the preconditioned residual vector (gc).

VMEC’s TimeStepControl tracks the minimum of the preconditioned residual (res0) and the physical residual (res1). If either grows by a factor of 1e4 after 10 steps, VMEC restores the last good state and reduces DELT by a factor of 1.03. vmec_jax mirrors this logic to reproduce the VMEC2000 iteration trace.

Multigrid interpolation (interp.f)

VMEC’s multigrid staging interpolates scaled coefficients between grids:

\[x_{\mathrm{scaled}} = x \cdot \mathrm{scalxc},\]

with odd-m extrapolation to the axis performed before interpolation. After linear interpolation on a uniform radial grid, coefficients are unscaled:

\[x_{\mathrm{new}} = \frac{x_{\mathrm{scaled,new}}}{\mathrm{scalxc}_{\mathrm{new}}}.\]

vmec_jax implements this exact pipeline so that stage-to-stage coefficient transfer matches VMEC2000.

Pressure and beta

VMEC reports thermal and magnetic energy scalars in wout:

\[W_B = \frac{1}{(2\pi)^2}\int \frac{B^2}{2}\,dV, \qquad W_P = \frac{1}{(2\pi)^2}\int p\,dV.\]

The total volume-averaged beta is computed by VMEC as:

\[\beta_{\mathrm{total}} = \frac{W_P}{W_B}.\]

vmec_jax follows the same normalization when emitting wout files.

Geometric embedding and Jacobian

Surfaces are represented in cylindrical coordinates by Fourier series:

\[R(s,\theta,\zeta) = \sum_{m,n}\left[ R_{mn}^c(s)\cos(m\theta-n\zeta) + R_{mn}^s(s)\sin(m\theta-n\zeta) \right],\]
\[Z(s,\theta,\zeta) = \sum_{m,n}\left[ Z_{mn}^c(s)\cos(m\theta-n\zeta) + Z_{mn}^s(s)\sin(m\theta-n\zeta) \right].\]

We embed into Cartesian coordinates using the physical toroidal angle:

\[x = R\cos\phi_{\mathrm{phys}}, \qquad y = R\sin\phi_{\mathrm{phys}}, \qquad z = Z.\]

The Jacobian is:

\[\sqrt{g} = \mathbf{e}_s \cdot (\mathbf{e}_\theta \times \mathbf{e}_{\phi}).\]

VMEC’s sign convention enforces signgs*sqrtg > 0 away from the axis.

Jacobian sign check (tau)

VMEC evaluates an auxiliary Jacobian-like scalar \(\tau\) built from even/odd-m real-space derivatives (see VMEC2000 jacobian.f). In compact form,

\[\tau \equiv (R_u Z_s - R_s Z_u) + \text{(odd-m corrections in } \sqrt{s}\text{)}.\]

If \(\tau\) changes sign away from the axis, VMEC flags a bad Jacobian and restarts the iteration with a refined axis guess. vmec_jax reproduces the same parity split, half-mesh averaging, and sign check so that Jacobian-reset behavior matches VMEC2000.

Implementation mapping (vmec_jax)

Key modules that directly implement the equations above:

  • vmec_jax/geom.py: geometry, metric, Jacobian.

  • vmec_jax/vmec_bcovar.py: contravariant/covariant field components.

  • vmec_jax/vmec_forces.py: real-space force kernels (A,B,C blocks).

  • vmec_jax/vmec_tomnsp.py: VMEC-style Fourier transforms of forces.

  • vmec_jax/vmec_residue.py: VMEC scalar residuals (FSQR/FSQZ/FSQL).

  • vmec_jax/solve.py: fixed-boundary iteration loop (VMEC2000 parity path).

References (local)

  • vmecpp/docs/the_numerics_of_vmecpp.pdf (VMEC++ numerics notes; VMEC2000 conventions and formulas).

  • VMEC2000 source in STELLOPT/VMEC2000/Sources.