Driver API

GMRESWorkspace(n, mem)

Storage for gmres! on vectors of length n with at most mem Arnoldi vectors: the basis V (n × (mem+1)), the Givens-rotated Hessenberg matrix R (upper triangular after rotation), the rotations c/s, and the rotated right-hand side z.

NKWorkspace(n, krylov_max)

Buffers for newton_krylov! on unknowns of length n: the residual at the current and the trial iterate, the Newton step, the finite-difference probe, the best iterate seen, and a GMRESWorkspace of krylov_max vectors.

compute_entropy_gradient!(ws, dS, v_parts, w_parts, f_coeffs_buf, r_vec, L_vec, G_buf)

Entropy gradient with respect to every particle velocity,

\[\frac{\partial S_h}{\partial v_\alpha} = -w_\alpha G_\alpha ,\]

where $G_\alpha = \nabla (M^{-1} r)(v_\alpha)$ is the finite-element projection of $\nabla \log f_s$ evaluated at particle $\alpha$ and $w_\alpha$ is its weight.

The projection chain is $v \mapsto f_s$ (L2 projection), $f_s \mapsto r$, then $r \mapsto M^{-1} r$ through the cached mass-matrix factorisation ws.M_lu. Both the projection and the gradient evaluation go through L2PROJ_FN and COMPG_FN, so this routine runs on the GPU when those are repointed.

Writes dS in place; f_coeffs_buf, r_vec, L_vec and G_buf are scratch.

gmres!(x, A!, b, gw::GMRESWorkspace; rtol, atol = 0.0, itmax = size(gw.R, 1))

Solve $A x = b$ from $x_0 = 0$ by unrestarted GMRES, with $A$ applied matrix-free as A!(y, u) ($y \leftarrow A u$). Stops once $\lVert b - A x_k \rVert \le \texttt{atol} + \texttt{rtol}\,\lVert b \rVert$, on Arnoldi breakdown, or after itmax products (capped by the workspace size).

Returns (k, rnorm): the number of products with $A$ and the final residual norm, which GMRES knows without forming $b - A x_k$.

Every product is one call to A!, which for newton_krylov! is one residual evaluation, so k is the cost that matters.

landau_geff!(ws, G_eff, v_in, v0, w_parts, S0, v_mid, dv, dS_mid, f_buf, r_vec,
             L_vec, G_buf; use_gonzalez = true)

The Landau discrete gradient $\overline{\nabla} S$ of picard_map!, as G_eff = -(∇S(v_mid) + λ Δv) / w: the FEM entropy gradient at the midpoint of v0 and v_in plus the Gonzalez rank-one term (λ = 0 with use_gonzalez = false). Sets v_mid and dv; everything up to the collision sum, so step_defect! can re-evaluate it without the pair sum.

newton_krylov!(x, F!, nk::NKWorkspace; tol, max_evals, fd_h,
               eta_max = 0.9, eta_gamma = 0.9, armijo = 1e-4,
               max_backtrack = 8, verbose = false)

Drive $\lVert F(x) \rVert$ below tol by inexact Newton, starting from x and overwriting it. F!(y, x) writes $F(x)$ into y.

The Newton step

Each step solves $J(x_k)\,\delta = -F(x_k)$ with gmres!, never forming $J$: a Jacobian–vector product is the forward difference

\[J u \approx \frac{F(x + h u) - F(x)}{h}, \qquad h = \texttt{fd\_h}\,/\,\lVert u \rVert ,\]

so each GMRES iteration costs one residual evaluation and the probe $x + h u$ always moves by exactly fd_h. The step is absolute on purpose. It must be small against the scale on which $J$ changes, which can be far smaller than $\lVert x \rVert$, and large against the noise in $F$. For the Landau step a Taylor test puts the first limit near $10^{-5}$ along its stiff direction, while a step scaled by $\sqrt{\varepsilon}\,\lVert x \rVert$ comes out at $\approx 0.1$ in FP32 (see docs/src/newton_krylov.md).

Forcing term

The linear solve is only as tight as the Newton step needs. GMRES stops at relative residual $\eta_k$, chosen by Eisenstat–Walker (choice 2),

\[\eta_k = \gamma \left(\frac{\lVert F_k \rVert}{\lVert F_{k-1} \rVert}\right)^2 ,\]

safeguarded from dropping abruptly ($\eta_k \ge \gamma\,\eta_{k-1}^2$ while that exceeds $0.1$), capped at eta_max, and kept above $\tfrac12\,\texttt{tol}/\lVert F_k \rVert$ so the last step does not over-solve past the target. The first step uses $\eta_0 = \texttt{eta\_max}$.

Globalisation

The step is accepted along $x_k + \lambda\delta$ with $\lambda$ halved from 1 until the Armijo condition $\lVert F(x_k + \lambda\delta) \rVert \le (1 - \alpha\lambda(1 - \eta_k))\lVert F_k \rVert$ holds. If max_backtrack halvings do not get there, the residual has hit the noise floor of $F$ and the solve stops.

Returns

(n_evals, rnorm, n_newton, rnorm0, status) where n_evals counts every call to F! (the solver's cost), status is :converged, :stalled (line search failed) or :budget (max_evals reached). On :stalled/:budget, x and nk.F hold the best iterate seen. In every case nk.F holds $F(x)$.

picard_map!(ws, v_out, v_in, v0, w_parts, S0, dt, v_mid, dv, dS_mid, G_eff,
            dot_v_buf, f_buf, r_vec, L_vec, G_buf; use_gonzalez = true)

One Picard map of the implicit step. Given a trial end-of-step state $v_\text{in}$ it returns the next iterate

\[v_\text{out} = v_0 + \Delta t \, \dot{v}(v_\text{mid}, \overline{\nabla} S) , \qquad v_\text{mid} = \tfrac{1}{2}\,(v_0 + v_\text{in}) , \qquad \Delta v = v_\text{in} - v_0 ,\]

so a fixed point $v_\text{out} = v_\text{in}$ is a solution of the implicit step. step_anderson! is what drives it to that fixed point.

Discrete gradient

With use_gonzalez = true the entropy gradient carries the Gonzalez correction

\[\overline{\nabla} S = \nabla S(v_\text{mid}) + \lambda \, \Delta v , \qquad \lambda = \frac{S(v_\text{in}) - S_0 - \Delta v \cdot \nabla S(v_\text{mid})} {\lVert \Delta v \rVert^2} ,\]

with $S_0 = S(v_0)$ passed in as S0. The rank-one term is chosen so that the discrete chain rule

\[\overline{\nabla} S \cdot \Delta v = S(v_\text{in}) - S_0\]

holds exactly rather than to within truncation error, which is what makes the scheme's entropy production discrete-exact.

use_gonzalez = false drops $\lambda$ and leaves the plain implicit-midpoint rule $\overline{\nabla} S = \nabla S(v_\text{mid})$. That removes the $\lVert \Delta v \rVert^2$ denominator, which is ill-conditioned when the run starts at or near equilibrium ($\Delta v \to 0$); in both branches the denominator is guarded by a $10^{-30}$ floor.

Collision model

ws.p.collision_model selects how $\dot{v}$ is assembled from $\overline{\nabla} S$:

• :landau — the $O(N^2)$ pair sum through COLLISION_FN, on a finite-element entropy gradient.
• :lb — the $O(N)$ Lenard–Bernstein drift, built from compute_moments and compute_drift_multipliers so momentum and energy are conserved exactly. LB uses the direct clamped $\nabla f_s / f_s$ as its base gradient in both modes: the FEM-projected seed injects Gibbs ringing into low-density cells, which Landau's $f$-weighted mobility suppresses and LB's does not. The discrete-gradient identity holds for any base gradient, so Gonzalez only adds the rank-one $\lambda$ term on top.

Writes v_out in place; every other array argument is scratch.

step_anderson!(ws, v1, v0, w_parts, S0, dt, ...; m = 5, max_iter = 1000,
               tol = 1e-12, abs_floor = 1e-7, damping = 0.5, ...)
    -> (iterations, residual, n_restarts, initial_residual)

Drive picard_map! to its fixed point with Anderson acceleration. v1 enters as the initial guess and leaves holding the solution. With use_anderson = false this degrades to plain damped Picard.

Anderson acceleration

Writing the Picard map as $G$ and the residual as $r_k = G(v_k) - v_k$, the last $m$ differences are kept in $\Delta F$ and $\Delta G$. Each iteration solves the small least-squares problem

\[\gamma = \arg\min_\gamma \lVert r_k - \Delta F \gamma \rVert_2^2 , \qquad v_{k+1} = (1 - \beta)\, v_k + \beta \left( G(v_k) - \Delta G \gamma \right) ,\]

with damping $\beta$ blending the accelerated step back toward the raw Picard update. The normal equations are regularised with $\lambda^2 = \texttt{reg\_factor} \cdot \operatorname{mean}(\operatorname{diag}(\Delta F^\top \Delta F))$ so the $m \times m$ system stays solvable when $\Delta F$ is near rank-deficient.

The window is a ring buffer: the newest difference overwrites the oldest column rather than shifting the whole window left, which is valid because the least-squares problem is invariant under a common column permutation of $\Delta F$ and $\Delta G$. See The Anderson window update for the measured cost.

Exit conditions

Any one of these ends the solve successfully. With exit_picard_step = true (the default) v1 leaves holding $G(v)$ for the converged iterate $v$, or for the best one seen on a stagnation or cap exit. That final Picard update is what keeps the step conservative. With $m = \tfrac12(v_0 + v)$ and $F = v - G(v)$, the collision operator conserves momentum and energy at $m$, so

\[P(G(v)) - P(v_0) = 0 , \qquad E(G(v)) - E(v_0) = -\tfrac{\Delta t}{2} \textstyle\sum_\alpha w_\alpha F_\alpha \cdot \dot v_\alpha(m) ,\]

whereas returning $v$ itself leaves $\sum_\alpha w_\alpha F_\alpha$ and $\sum_\alpha w_\alpha m_\alpha \cdot F_\alpha$, first order in the residual. The returned residual is $\lVert F(v) \rVert$, not that of $G(v)$. Along the stiff direction of the Landau map the update amplifies the residual (a median of about 2x over 30 FP64 steps, scripts/exit_update_check.jl), which is the price of exact momentum conservation. exit_picard_step = false returns $v$ instead; an FP32 A/B showed no cost benefit from it, and 3–10x worse energy drift (docs/src/newton_krylov.md).

1. Residual. $\lVert r \rVert < \max(\texttt{tol} \cdot \lVert v \rVert, \texttt{abs\_floor})$. abs_floor caps how tight the solve is asked to be: below the numerical noise floor of the Picard map the extra iterations buy nothing. In FP32 the floor matters a great deal: abs_floor = 1e-8 cuts the iteration count by over 3x against the 1e-10 the presets carry, at no cost in accuracy.
2. Stagnation. Every stag_window iterations the best residual is compared with its value one window earlier; a relative drop below stag_rel_tol exits. This catches the late-step plateau where Anderson cannot push further. An iteration count that is an exact multiple of stag_window is the signature of a step that left by this route rather than by converging.
3. Iteration cap. max_iter, which warns.

Past damp_decay_start iterations without an exit, damping is multiplied by damp_decay_factor for a more conservative step, which stabilises stiff late-time maps. If restart_factor is finite, the window is cleared whenever the residual exceeds that multiple of the best residual so far.

step_defect!(ws, v1, v0, w_parts, S0, dt, v_mid, dv, dS_mid, G_eff, dot_v_buf,
             f_buf, r_vec, L_vec, G_buf, Gv, r_curr, r_prev, Gv_prev, v_old, ΔF, ΔG,
             A, Gk, Φu; m = 8, max_iter = 2000, tol = 1e-12, abs_floor = 1e-7,
             damping = 0.7, eta = 0.05, max_inner = 60, stag_window = 5,
             stag_rel_tol = 0.1, use_gonzalez = true, verbose = false)
    -> (outer_iterations, residual, inner_iterations, initial_residual)

Solve the Landau implicit step of step_anderson! by defect correction around a frozen metric. The collision velocity splits as

\[\dot v_\gamma = \sum_\alpha w_\alpha U(v_\gamma - v_\alpha)(G_\alpha - G_\gamma) = B_\gamma - A_\gamma G_\gamma , \qquad A_\gamma = \sum_\alpha w_\alpha U(v_\gamma - v_\alpha) ,\]

where the pair sums $A$ and $B$ change slowly with the particle positions while the entropy gradient $G$ (an $O(N)$ projection) changes fast. Outer iteration $k$ evaluates the full Picard map once, $\mathcal{G}(v_k)$, together with $A_k$ in the same pair pass (COLLMETRIC_FN), and then solves the cheap map

\[\Phi_k(u) = \mathcal{G}(v_k) - \Delta t\, A_k \bigl(G_\text{eff}(u) - G_\text{eff}(v_k)\bigr)\]

by Anderson to $\lVert \Phi_k(u) - u \rVert < \max(\eta \lVert \mathcal{G}(v_k) - v_k \rVert, \texttt{abs\_floor}/2)$. $\Phi_k(v_k) = \mathcal{G}(v_k)$, so a fixed point of the outer iteration is a fixed point of $\mathcal{G}$: the root is the same as step_anderson!'s. Each $\Phi_k$ evaluation needs $G_\text{eff}$ (landau_geff!) but no pair sum.

The stopping rule is step_anderson!'s on the outer residual $\lVert \mathcal{G}(v) - v \rVert$, and the exit returns $\mathcal{G}(v)$ computed with the antisymmetric pair sum, so momentum and energy are conserved exactly as there. Every stag_window outer iterations, a best residual that fell by less than stag_rel_tol ends the solve with the best $\mathcal{G}(v)$. max_iter caps the outer iterations.

Measured (N = 40k, FP64, RTX 3080 probe): 6–11 outer iterations per step against 10–120 Anderson iterations, each $\Phi_k$ costing about 2 % of a full map. In FP32 the full map is cheap and this does not pay.

step_newton!(ws, v1, v0, w_parts, S0, dt, ..., Gv, nk::NKWorkspace;
             max_iter = 1000, tol = 1e-12, abs_floor = 1e-7, fd_h = 1e-5,
             fd_rel = 0.0, eta_max = 0.9, verbose = false,
             use_gonzalez = true, exit_picard_step = true)

Solve the same implicit step as step_anderson! by Jacobian-free Newton–Krylov (newton_krylov!) on the residual of the Picard map,

\[F(v) = v - \mathcal{G}(v), \qquad \mathcal{G} = \texttt{picard\_map!} ,\]

from the predictor already in v1. The Jacobian $I - \partial\mathcal{G}/\partial v$ is $I - O(\Delta t)$, so GMRES needs few iterations per Newton step.

The stopping rule and the cost unit match step_anderson!, so the two are directly comparable: the target is $\max(\texttt{tol}\cdot\lVert v \rVert, \texttt{abs\_floor})$, with $\lVert v \rVert$ taken once at the predictor rather than at every iterate (a negligible difference: it changes by about $10^{-3}$ relative within a step), max_iter caps the number of Picard-map evaluations, and the returned count is the number of evaluations — each finite-difference Jacobian product and each line-search trial is one. On exit v1 takes the final Picard update $\mathcal{G}(v) = v - F(v)$, as in step_anderson!, which restores momentum and energy conservation; exit_picard_step = false returns $v$ itself.

The finite-difference step is the absolute fd_h (see newton_krylov!). fd_rel > 0 overrides it with the earlier relative scaling $h = \texttt{fd\_rel}\cdot\max(\lVert v \rVert, 1)$, for reproducing old runs.

Returns (n_evals, ‖F‖, n_newton, ‖F₀‖); the third slot is the restart count in step_anderson!.

COLLISION_FN, L2PROJ_FN, COMPG_FN, LOGGRAD_FN

Swappable implementations of the four hot-loop kernels, held in Refs so the backend can be chosen at startup rather than through the type system.

They default to the CPU routines from functions.jl. Passing --use_gpu=true makes main load collision_gpu.jl and projection_gpu.jl and repoint these at the CUDA kernels: the $O(N^2)$ Landau pair sum (COLLISION_FN, FP64 or FP32), the particle-to-spline L2 projection (L2PROJ_FN), the spline gradient gather (COMPG_FN) and, for Lenard–Bernstein, the direct log-density gradient (LOGGRAD_FN).

Every call site goes through Base.invokelatest so the swap is visible to code that was already compiled, which a plain call would miss under world-age.

COLLISION_FN, L2PROJ_FN, COMPG_FN, LOGGRAD_FN

Swappable implementations of the four hot-loop kernels, held in Refs so the backend can be chosen at startup rather than through the type system.

They default to the CPU routines from functions.jl. Passing --use_gpu=true makes main load collision_gpu.jl and projection_gpu.jl and repoint these at the CUDA kernels: the $O(N^2)$ Landau pair sum (COLLISION_FN, FP64 or FP32), the particle-to-spline L2 projection (L2PROJ_FN), the spline gradient gather (COMPG_FN) and, for Lenard–Bernstein, the direct log-density gradient (LOGGRAD_FN).

Every call site goes through Base.invokelatest so the swap is visible to code that was already compiled, which a plain call would miss under world-age.

COLLISION_FN, L2PROJ_FN, COMPG_FN, LOGGRAD_FN

Swappable implementations of the four hot-loop kernels, held in Refs so the backend can be chosen at startup rather than through the type system.

They default to the CPU routines from functions.jl. Passing --use_gpu=true makes main load collision_gpu.jl and projection_gpu.jl and repoint these at the CUDA kernels: the $O(N^2)$ Landau pair sum (COLLISION_FN, FP64 or FP32), the particle-to-spline L2 projection (L2PROJ_FN), the spline gradient gather (COMPG_FN) and, for Lenard–Bernstein, the direct log-density gradient (LOGGRAD_FN).

Every call site goes through Base.invokelatest so the swap is visible to code that was already compiled, which a plain call would miss under world-age.

COLLISION_FN, L2PROJ_FN, COMPG_FN, LOGGRAD_FN

Swappable implementations of the four hot-loop kernels, held in Refs so the backend can be chosen at startup rather than through the type system.

They default to the CPU routines from functions.jl. Passing --use_gpu=true makes main load collision_gpu.jl and projection_gpu.jl and repoint these at the CUDA kernels: the $O(N^2)$ Landau pair sum (COLLISION_FN, FP64 or FP32), the particle-to-spline L2 projection (L2PROJ_FN), the spline gradient gather (COMPG_FN) and, for Lenard–Bernstein, the direct log-density gradient (LOGGRAD_FN).

Every call site goes through Base.invokelatest so the swap is visible to code that was already compiled, which a plain call would miss under world-age.

COLLISION_FN, L2PROJ_FN, COMPG_FN, LOGGRAD_FN

Swappable implementations of the four hot-loop kernels, held in Refs so the backend can be chosen at startup rather than through the type system.

They default to the CPU routines from functions.jl. Passing --use_gpu=true makes main load collision_gpu.jl and projection_gpu.jl and repoint these at the CUDA kernels: the $O(N^2)$ Landau pair sum (COLLISION_FN, FP64 or FP32), the particle-to-spline L2 projection (L2PROJ_FN), the spline gradient gather (COMPG_FN) and, for Lenard–Bernstein, the direct log-density gradient (LOGGRAD_FN).

Every call site goes through Base.invokelatest so the swap is visible to code that was already compiled, which a plain call would miss under world-age.

show(io, ::MIME"text/plain", p::SimParameters)

Print p in keyword-constructor form, one field per line. The output is valid Julia, so PARAMS = <output> is itself a preset that reproduces the run.

build_workspace(p::SimParameters)

Construct geometry, function space, mass matrix, quadrature, and all preallocated scratch from the configuration p. Each velocity dimension uses its own breakpoint vector supplied by p.bp1, p.bp2 (anisotropic, possibly non-uniform). The breakpoints are aliased into the workspace without copying.

compute_collision_metric!(ws, dot_v, A, v_parts, w_parts, G)

compute_collision! plus, in the same pair loop, the metric $A_\gamma = \sum_\alpha w_\alpha U(v_\gamma - v_\alpha)$ stored as A[γ, :] = (A₁₁, A₁₂, A₂₂). Then $\dot v_\gamma = B_\gamma - A_\gamma G_\gamma$ with $B$ the $U$-weighted sum of the other particles' $G$: what step_defect! freezes. dot_v is accumulated in the same antisymmetric form as compute_collision!, so it is bit-identical to it.

parse_overrides(p::SimParameters, args)

Apply every --key=value token in args on top of p, returning a new SimParameters. Unknown keys raise. Empty args returns p unchanged.

print_summary(p::SimParameters)

One-shot human-readable dump of the active configuration. Used at the top of each run so the .log captures exactly what was run.