argminax: optimisation with uncertainty for JAX

argminax is the stack’s optimisation library with uncertainty at the optimum: it finds x̂ = argmin J(x), and then says how sharp that minimum is (posterior covariance, averaging kernel, degrees of freedom) and how x̂ depends on what went in. It sits at Core tier 1 on gaussx, extends optimistix with the stack’s solvers and adjoints, and takes over the curvature code now rebuilt in seven places. Status: proposed, at design-sketch depth.

At a glance¶

argminax sits one tier above gaussx, whose single steps it iterates, so every library that minimises can import it. This page is the design; core types, algorithms and the code survey are in argminax: design details.

The problem¶

A methane retrieval is not finished when the optimiser stops. radtrax’s 3D-Var (retrieval/var3d) finds the column enhancement that best fits a hyperspectral pixel, and then needs three more things at that minimum: how uncertain the answer is, how much of it came from the measurement rather than the prior, and how its gradient flows back to the plume model that called it. Those three are the curvature at the minimum, and every library that optimises rebuilds them:

The solvers mostly come from optimistix already. What each library rewrites is what optimistix leaves out: Gauss–Newton curvature as a structured operator, its inverse as a posterior covariance, the degrees of freedom of a retrieval, and adjoints for differentiating through a minimiser that stops early. Seven copies of that machinery, across GeoModels and Core, with plumax and radtrax each holding the same 3D-Var solvers.

The cost, its minimum, and the curvature there¶

A retrieval minimises a cost that balances the measurement against the prior. Here y is the measured radiance, F the forward model, x the state (the methane column), xₐ the prior state, Sε the measurement-error covariance and Sₐ the prior covariance:

Three jobs share the word optimisation, and only one belongs here:

At the minimum x̂, with K = ∂F/∂x the Jacobian there, the Gauss–Newton curvature gives the posterior covariance Ŝ and the averaging kernel A:

A says how much a change in the true state moves the retrieved one: A = I means the measurement decides everything, A = 0 means the prior does. Its trace is the retrieval’s degrees of freedom for signal. The same Ŝ is the inverse Gauss–Newton Hessian of 4D-Var, the Laplace approximation in pyrox and the precision in gaussx; only the names differ.

The last piece is the derivative of x̂ with respect to what went in (the prior, the parameters of F). At an exact minimum it comes from the implicit function theorem; for a solver stopped early, as in learned 4D-Var, it comes from differentiating the last k iterations. Both are adjoint strategies of the minimiser.

Modules and boundaries¶

Four modules: the solvers optimistix lacks, the curvature at the minimum, the retrieval diagnostics built on it, and the adjoints of a minimiser. argminax extends optimistix as gaussx extends lineax: its solvers are optimistix solvers and its adjoints are optimistix adjoints, so optx.minimise and optx.least_squares keep working.

What argminax owns, and what it hands off¶

The test for every name: does it find a minimum, or say something about the curvature at one? Then argminax owns it.

What moves in¶

In use¶

The methane retrieval end to end: solve for the column, then read the uncertainty, the averaging kernel and the degrees of freedom off the curvature at the minimum, and differentiate the answer with respect to the prior. Every argminax name is proposed; the solve itself is ordinary optimistix.

# Shapes: W spectral channels; L layers in the state; Sε is (W, W), Sₐ is (L, L)

import argminax as amx
import jax
import optimistix as optx


# user
# r(x) = (Sε^(-1/2) (y − F(x)), Sₐ^(-1/2) (x − xₐ)), so J(x) = ½ ‖r(x)‖²
def residuals(state, args): ...  # (L,) → (W + L,)


# proposed (argminax)
# Gauss–Newton with an inner CG on (KᵀSε⁻¹K + Sₐ⁻¹) δ = −∇J
solver = amx.solve.GaussNewtonCG(rtol=1e-6, atol=1e-6)
# not pinned (optimistix)
# d x̂ / d(xₐ, θ) by the implicit function theorem at x̂
adjoint = optx.ImplicitAdjoint()
# not pinned (optimistix)
# x̂ = argmin J(x)
solution = optx.least_squares(  # (L,) → (L,)
    residuals, solver, prior_state, args=problem, adjoint=adjoint
)
retrieved = solution.value  # → (L,)

# --- curvature at the minimum ---
# proposed (argminax)
# Ŝ = (KᵀSε⁻¹K + Sₐ⁻¹)⁻¹, a gaussx operator, never formed densely
posterior = amx.curvature.gauss_newton(residuals, retrieved, problem)  # → (L, L) op
# proposed (argminax)
# A = I − Ŝ Sₐ⁻¹
kernel = amx.diagnose.averaging_kernel(posterior, prior_covariance)  # → (L, L) op
# proposed (argminax)
# d_s = tr(A), the degrees of freedom for signal
dofs = amx.diagnose.degrees_of_freedom(kernel)  # → ()

# --- the retrieval inside a larger model ---
# user: retrieve(xₐ) wraps the least_squares call above
# x̂ is differentiable: ∂x̂/∂xₐ flows through the implicit adjoint
sensitivity = jax.jacrev(retrieve)(prior_state)  # (L,) → (L, L)

A plume inversion calls retrieve per pixel and differentiates through it; the implicit adjoint costs one more linear solve with the same curvature, not a pass through every iteration. A learned 4D-Var that stops after a fixed number of steps swaps in amx.adjoint.KStep(k=3).

Integration across the stack¶

Six libraries would import argminax. For radtrax, plumax and vardax the solve call barely changes, because they already use optimistix: what changes is that the curvature, diagnostics and adjoints come from one place.

vardax’s adjoint strategies are the clearest move: they are already optimistix.AbstractAdjoint subclasses, so for vardax only the import changes. They are not yet usable elsewhere: their apply() raises NotImplementedError, and they work only inside vardax’s own unrolled solver. Wave 1 implements apply, so they work under any optimistix solve.

# Today:
# today (vardax)
from vardax.adjoints import KStepAdjoint

# After:
# proposed (argminax)
from argminax.adjoint import KStep

vardax keeps a deprecated alias for one release, and FourDVarNet1D(solver_adjoint=...) call sites do not change.

Decisions and open questions¶

Five decided, three open. Each entry gives the position first, then the trade-off.

1. Name: argminax¶

Decided (5 Nov 2026). argminax, free on PyPI when checked. It extends optimistix, as gaussx extends lineax.

2. Scope: solve and curvature, not training¶

Decided (5 Nov 2026). argminax owns deterministic solves to a tolerance and what the curvature at the minimum says. Stochastic training over batches stays with optax and optax_bayes; sampling and variational inference stay with NumPyro and pyrox.

3. Layer: Core tier 1¶

Decided (5 Nov 2026, on the layer map). argminax imports gaussx (Core tier 0) and extends optimistix. vardax, pyrox-gp, pyrox-lgm, pyrox-nn, radtrax and plumax sit above it.

4. gaussx keeps one step; argminax owns the loop¶

Decided (5 Nov 2026; revised 6 Nov 2026). gaussx keeps newton_update, gauss_newton_precision and its Chandrupatla root finder, because they are single-step algebra a tier below. laplace_mode, the loop that iterates them, moves to argminax whole, entry point included, and its callers (pyrox-gp, pyrox-lgm) import it from there. gaussx never calls argminax: a laplace_mode left in gaussx that forwarded to argminax would be an edge from Core tier 0 up to tier 1, and the repository fixes such an edge by moving code, never by a re-export (AGENTS.md, “Boundaries”). gaussx’s own laplace_mode keeps its current loop for one release behind a DeprecationWarning that names argminax.solve.laplace_mode, then goes.

5. vardax’s adjoint strategies move here¶

Decided (5 Nov 2026). KStepAdjoint and OneStepAdjoint are optimistix.AbstractAdjoint subclasses: they differentiate through a minimiser, so they are argminax’s. Today their apply() raises NotImplementedError and they work only inside vardax’s unrolled solver; the move implements apply, so they work under any optx.minimise or optx.least_squares call. Backward sweeps through a forward run stay integrix’s (integrix decision 7).

6. Where radtrax’s 3D-Var solvers live¶

Open. run_gauss_newton and run_dual_psas are general (any sum of squared residuals; any linear observation operator), so they fit amx.solve. run_lbfgs is a thin wrapper over optimistix and may simply go. The cost, control transforms and whitening stay in radtrax. Either way, plumax’s copy in assimilation/solve.py goes.

7. Positivity and bounds¶

Leaning: by reparametrisation. Methane enhancements and emission rates are positive. A log or softplus transform keeps the solve unconstrained and the curvature well defined; the posterior is then reported in the transformed space, with a note on how to map it back. Projected methods stay out until a case needs them.

8. Diagnostics: argminax or radtrax?¶

Open. The averaging kernel, degrees of freedom and information content come from atmospheric retrieval Rodgers, 2000, but they are properties of any linear-Gaussian posterior: the same A describes a 4D-Var analysis or a GP posterior. Leaning: argminax, with the retrieval vocabulary in its docstrings.

Roadmap¶

Five waves, each a GitHub issue epic in the argminax repository once it exists. The adjoints and curvature come first: they already exist in vardax and need a home, plus the apply method they lack, and every later wave uses them.

Wave 4 is one issue per consumer, filed in that consumer’s repository and linked to the argminax epic. pyrox-nn’s issue follows pyrox#332, which moves last_layer_laplace out of the pyrox core.

Sources¶

• Optimistix documentation: the solver and adjoint protocols argminax extends.

• Rodgers (2000): the averaging kernel, degrees of freedom and information content.

• Bolte et al. (2023): one-step differentiation of iterative algorithms, the basis of vardax’s k-step adjoint, as its docstring cites.

• The GeoStack layer map: argminax’s layer, tier, imports and moves.

• integrix: the adjoint boundary between the two.

• pyrox#332: last_layer_laplace out of the pyrox core.

• Code read on 5 Nov 2026 at geostack’s submodule pins: radtrax retrieval/var3d/solve.py and diagnostics.py; plumax les_fvm/fourdvar.py and assimilation/solve.py; vardax adjoints/ and posterior/; pyrox-gp _inference_nongauss.py; pyrox-lgm _inla.py; gaussx _inference/.

• pyrox last_layer_laplace read on pyrox main at 2ac68a5; geostack’s pyrox pin (94eef13) does not have it yet.