Method overview

dkx implements a radially local neoclassical drift-kinetic solve stack with structured velocity-space discretization, matrix-free operator application, and CPU/GPU-capable solver orchestration.

The code solves for the non-adiabatic correction to the distribution function on a single flux surface. For the model and literature background, see Physics model and equations and Physics reference: the radially local drift-kinetic model. This page focuses on the discretization and implementation choices that turn that model into an efficient computational workflow.

Example Boozer field magnitude BHat(θ, ζ)

Example v3 geometryScheme=11/12 Boozer field magnitude \(\hat B(\theta,\zeta)\) on a discrete grid.

Discretization (v3)

SFINCS v3 uses a structured discretization that is well-suited to a JAX port:

  • Angles: periodic grids in \(\theta\) and \(\zeta\) with finite-difference derivative matrices.

  • Speed: a polynomial/Stieltjes grid in \(x\) (a normalized speed-like coordinate), with quadrature weights used in moments/constraints.

  • Pitch angle: a Legendre-mode expansion in \(\xi = v_\parallel / v\).

The primary unknown (in many modes) can be viewed as a tensor

\[f = f(s, x, L, \theta, \zeta),\]

where \(s\) is the species index and \(L\) is the Legendre index.

Geometry harmonics and resolution (v3 detail)

For Boozer-coordinate geometries that are represented by Fourier harmonics (e.g. geometryScheme=11/12), SFINCS v3 truncates the harmonic tables to modes representable on the discrete \((\theta,\zeta)\) grid. Concretely, for a grid of sizes Ntheta and Nzeta, the v3 geometry builder includes only modes satisfying

\[0 \le m \le \lfloor N_\theta/2 \rfloor, \qquad |n| \le \lfloor N_\zeta/2 \rfloor,\]

with additional Nyquist exclusions for sine components.

dkx matches this behavior when evaluating .bc geometries so that the resulting arrays are consistent with the v3 Jacobian assembly at a given resolution.

Implemented operator terms

The drift-kinetic operator is assembled from the following parity-tested building blocks.

Collisionless streaming + mirror (ΔL = ±1)

The basic collisionless dynamics couple neighboring Legendre modes \(L \leftrightarrow L\pm 1\) through (i) streaming along the field line and (ii) the mirror force. These terms are parity-tested against frozen PETSc binaries.

ExB drift terms (ΔL = 0)

SFINCS v3 includes ExB advection in the angular directions. In the Fortran code (populateMatrix.F90), these terms are assembled as dense derivative-matrix applications in \(\theta\) and \(\zeta\), with a coefficient proportional to \(\partial_{\psi}\hat\Phi\).

For the default (non-DKES) form used in most parity tests, the coefficients are:

\[F_{E\times B,\theta}(\theta,\zeta) = \frac{\alpha\,\Delta\,\partial_{\psi}\hat\Phi}{2} \;\frac{\hat D\;\hat B_{\zeta}}{\hat B^2}, \qquad F_{E\times B,\zeta}(\theta,\zeta) = -\frac{\alpha\,\Delta\,\partial_{\psi}\hat\Phi}{2} \;\frac{\hat D\;\hat B_{\theta}}{\hat B^2}.\]

In geometryScheme=4, \(\hat B_{\theta} = \hat I\) is zero in the default W7-X parameter set, so the \(\partial/\partial\zeta\) ExB term vanishes.

Non-standard Er term in xiDot (ΔL = ±2)

SFINCS v3 contains an additional, non-standard \(\partial/\partial\xi\) term associated with the radial electric field. In the Legendre basis, this term has a diagonal-in-\(L\) piece and couples \(L \leftrightarrow L\pm 2\).

In the Fortran code (populateMatrix.F90), the coefficient for this term is

\[F_{\xi}(\theta,\zeta) = \frac{\alpha\,\Delta\,\partial_{\psi}\hat\Phi}{4\,\hat B^3} \;\hat D\;\Big( \hat B_{\zeta}\,\partial_{\theta}\hat B - \hat B_{\theta}\,\partial_{\zeta}\hat B \Big),\]

where the hats denote v3-normalized quantities.

In the geometryScheme=4 model used in tests, \(\hat B_{\theta}\) is constant (and is zero in the default W7-X parameter set), so the expression simplifies.

Collisionless Er xDot term (x-coupling and ΔL = ±2)

When includeXDotTerm = .true., v3 includes a collisionless radial-derivative term in the kinetic equation, discretized as a dense differentiation matrix in the \(x\) coordinate.

The v3 implementation includes:

  • Dense x-matvec using \(x\,\partial/\partial x\) (with optional upwinding schemes).

  • Legendre couplings with \(L \leftrightarrow L\pm 2\) (and a diagonal-in-\(L\) piece).

dkx implements the default xDotDerivativeScheme = 0, i.e. the same polynomial-grid differentiation matrix is used for both upwind directions.

Magnetic drift terms (ΔL = 0 and ΔL = ±2)

Note

Every magneticDriftScheme from 0 (off, the default) through 9 is built by the consolidated dkx.drift_kinetic.KineticOperator. When magneticDriftScheme > 0 the coefficient forms below are assembled and applied inside KineticOperator.apply_f. Because the drift term couples \(L \leftrightarrow L\pm2\) it breaks the block-tridiagonal-in-\(L\) structure, so KineticOperator.to_block_tridiagonal declines drift decks and the tier-2 recycled-Krylov solver owns them (Numerics and algorithms). The full versus DKES trajectory distinction is carried separately by the \(E_r\) energy/pitch terms and the useDKESExBDrift switch (Physics reference: the radially local drift-kinetic model).

SFINCS v3 includes magnetic-drift advection in the angular directions, plus an associated non-standard \(\partial/\partial\xi\) term.

For the v3 Jacobian assembly (populateMatrix.F90), the d/dtheta and d/dzeta drift terms take the form

\[\mathcal{L}_{\mathrm{md},\theta}[f] = F_{\mathrm{md}}(\theta,\zeta)\;\partial_{\theta} f \quad\text{(with Legendre couplings in } L\text{)},\]
\[\mathcal{L}_{\mathrm{md},\zeta}[f] = F_{\mathrm{md}}(\theta,\zeta)\;\partial_{\zeta} f \quad\text{(with Legendre couplings in } L\text{)}.\]

Each term includes a diagonal-in-\(L\) piece and \(L\leftrightarrow L\pm 2\) couplings. In the Fortran code, the shared prefactor is

\[\mathrm{factor}(\theta,\zeta,x) = \frac{\Delta\,\hat T\,\hat D\,x^2}{2\,Z\,\hat B^3}.\]

For magneticDriftScheme = 1 (the first scheme ported), the v3 code defines the geometric factors (using v3-normalized quantities)

\[g_1^{(\theta)} = \hat B_{\zeta}\,\partial_{\psi}\hat B - \hat B_{\psi}\,\partial_{\zeta}\hat B, \qquad g_2^{(\theta)} = 2\,\hat B\,\big(\partial_{\zeta}\hat B_{\psi} - \partial_{\psi}\hat B_{\zeta}\big),\]

and (for the zeta drift term as implemented in dkx parity fixtures)

\[g_1^{(\zeta)} = \hat B_{\psi}\,\partial_{\theta}\hat B - \hat B_{\theta}\,\partial_{\psi}\hat B, \qquad g_2^{(\zeta)} = 2\,\hat B\,\big(\partial_{\psi}\hat B_{\theta} - \partial_{\theta}\hat B_{\psi}\big).\]

To stabilize the drift advection, v3 supports upwinded angular derivative matrices (ddtheta_magneticDrift_plus/minus and ddzeta_magneticDrift_plus/minus). When magneticDriftDerivativeScheme != 0, v3 selects the upwind direction at each \((\theta,\zeta)\) grid point based on

\[\mathrm{use\_plus}(\theta,\zeta) = \big(g_1\,\hat D(1,1)/Z\big) > 0.\]

This upwind selection is parity-tested against frozen PETSc binaries for a geometryScheme=11 fixture; it is active whenever magneticDriftScheme > 0 (the default magneticDriftScheme = 0 leaves the magnetic-drift terms off).

Collision operators (PAS, FP, and improved Sugama)

SFINCS v3 supports two collision models controlled by collisionOperator, and dkx adds a third:

  • collisionOperator = 1: pure pitch-angle scattering (PAS) in the Legendre basis.

  • collisionOperator = 0: full linearized Fokker-Planck operator (Landau form), implemented via Rosenbluth potentials (and dense coupling in the speed coordinate \(x\)).

  • collisionOperator = 3: the momentum- and energy-conserving improved Sugama model operator (a dkx research extension beyond Fortran v3); it shares the Fokker-Planck test-particle kernels and dense per-\(L\) speed-block structure.

For both models, the operator is diagonal in \((\theta,\zeta)\) and diagonal in the Legendre index \(L\). For a fixed \(L\), the action can be written as a dense x-space matrix-vector product that couples species:

\[(\mathcal{C} f)_{a, i, L, \theta, \zeta} = \sum_b\sum_j \mathsf{C}^{(L)}_{a b, i j}\; f_{b, j, L, \theta, \zeta}.\]

In the v3 matrix assembly, the overall normalization is applied via nu_n (see populateMatrix.F90). On the canonical stack this model is dkx.collisions (the pitch-angle-scattering and full Fokker–Planck operators), applied inside KineticOperator.apply_f and parity-tested by comparing a full F-block matvec against a frozen PETSc Jacobian for the v3 example quick_2species_FPCollisions_noEr.

Poloidally varying collisions (Phi1 in the collision operator)

When includePhi1InCollisionOperator = .true. (and includePhi1InKineticEquation = .true.), v3 modifies the collision operator coefficients through a poloidally varying effective density

\[n_a^{\mathrm{pol}}(\theta,\zeta) = n_a \exp\left(-\frac{Z_a \alpha}{T_a}\,\Phi_1(\theta,\zeta)\right),\]

so that the Fokker–Planck operator remains diagonal in \((\theta,\zeta)\) but is no longer uniform on the flux surface.

On the canonical stack this poloidally varying collision operator is applied inside dkx.drift_kinetic.KineticOperator.apply_f(), which rescales the dkx.collisions Fokker–Planck blocks by the poloidal density factor when includePhi1InCollisionOperator = .true.. It is parity-tested against a frozen v3 PETSc matrix for the fixture fp_1species_FPCollisions_noEr_tiny_withPhi1_inCollision. For derivations and implementation details, see the upstream SFINCS technical notes cited from Upstream SFINCS sources and primary literature.

Why JAX?

Using JAX enables:

  • JIT compilation (CPU/GPU) of the operator application and solver kernels.

  • Automatic differentiation through geometry, collision operators, and eventually the full kinetic solve (useful for sensitivity studies and gradient-based optimization).

  • Matrix-free linear algebra: express the v3 Jacobian as a matvec rather than assembling sparse matrices, enabling scalable iterative solvers.

  • An ecosystem of tools that become natural once the compute graph is differentiable:

    • optax for gradient-based optimization loops (calibration, inverse problems).

Code design strategy

The code is organized around a small number of principles:

  1. keep the physics terms explicit and inspectable,

  2. separate operator construction from solve-policy decisions,

  3. use JAX where it materially helps performance or differentiability,

  4. allow bounded host-side rescues when they are the practical production solution,

  5. validate every numerical change with targeted tests and audited case runs.

Matrix-free residual and Jacobian application

For iterative solvers and implicit differentiation, it is useful to work with a residual function rather than an assembled sparse matrix. For the linear kinetic block this residual is

\[r(x) = A x - b,\]

where \(A\) is represented by a matrix-free matvec and \(b\) is a right-hand side.

On the canonical stack the residual and Jacobian interface is the consolidated dkx.drift_kinetic.KineticOperator:

  • KineticOperator.apply(v) is the matrix-free matvec \(Av\), and KineticOperator.rhs() is the right-hand side \(b\).

  • For the linear kinetic block the Jacobian matvec is identical to the operator matvec; for the nonlinear \(\Phi_1\) residual residual_phi1(x), jax.jvp provides an efficient Jacobian-vector product without ever forming a dense or sparse Jacobian matrix.

Linear solvers and preconditioning

SFINCS v3 systems are large, stiff, and often ill-conditioned due to collision operators, constraint nullspaces, and mixed dense/sparse structure. The canonical stack keeps the operator matrix-free and applies a three-tier auto policy (dkx.solve.solve()) described in full in Numerics and algorithms:

  • Tier 1 — a structured direct solve (block-tridiagonal Legendre elimination) for the DKES-trajectory / pitch-angle-scattering family;

  • Tier 2 — matrix-free recycled Krylov (FGMRES + GCROT) right-preconditioned by an exact tier-1 solve of a SFINCS-simplified coarse operator (the Fortran preconditionerOptions idiom);

  • Tier 3 — a host sparse-direct (SuperLU) escape hatch.

The structured factorizations, recycled Krylov, and implicit-solve wrappers are provided by the solvax library.

Block-tridiagonal elimination for the Legendre chain

The tier-1 solve exploits the block-tridiagonal-in-\(L\) structure of the monoenergetic operator. For diagonal blocks \(D_k\) and neighbour couplings \(L_k\), \(U_k\),

\[\begin{split}A = \begin{bmatrix} D_0 & U_0 \\ L_0 & D_1 & U_1 \\ & \ddots & \ddots & \ddots \\ & & L_{n-2} & D_{n-1} \end{bmatrix},\end{split}\]

the block-Schur (block-Thomas) recursion is

\[S_0 = D_0, \qquad C_0 = S_0^{-1} U_0,\]
\[S_k = D_k - L_{k-1} C_{k-1}, \qquad C_k = S_k^{-1} U_k \quad (k \ge 1),\]

and repeated solves need only forward/backward substitution:

\[y_0 = S_0^{-1} b_0, \qquad y_k = S_k^{-1}(b_k - L_{k-1} y_{k-1}), \qquad x_k = y_k - C_k x_{k+1}.\]

On the canonical stack the blocks come from dkx.drift_kinetic.KineticOperator.to_block_tridiagonal(), and the factor/solve is the solvax block-Thomas kernel with a truncated-storage back-substitution that retains only the low-\(L\) blocks the right-hand side and moments touch. This is the Legendre block-tridiagonal method of Escoto (2025) (see Numerics and algorithms).