spectrax.midpoint_solver

Custom Diffrax solver: implicit midpoint with a Newton–GMRES nonlinear solve.

Classes

ImplicitMidpoint

Implicit midpoint ODE solver using a JAX-compiled Newton–GMRES iteration.

Functions

_newton_gmres(F_fn, y0, y_init, rtol, atol, max_iters)

Solve F(y)=0 using Newton iterations with GMRES linear solves.

Module Contents

class spectrax.midpoint_solver.ImplicitMidpoint

Bases: diffrax.AbstractSolver

Implicit midpoint ODE solver using a JAX-compiled Newton–GMRES iteration.

This solver implements the implicit midpoint rule:

y_{n+1} = y_n + Δt * f(t_{n+1}, (y_n + y_{n+1}) / 2)

The nonlinear equation for y_{n+1} is solved with Newton iterations, where each linearized step is solved via GMRES using JAX’s linearization.

rtol: float = 1e-06

atol: float = 1e-08

max_iters: int = 300

term_structure

interpolation_cls

order(terms)

init(terms, t0, t1, y0, args)

func(terms, t0, y0, args)

step(terms, t0, t1, y0, args, solver_state, made_jump)

spectrax.midpoint_solver._newton_gmres(F_fn, y0, y_init, rtol, atol, max_iters)

Solve F(y)=0 using Newton iterations with GMRES linear solves.

Notes

• The Jacobian-vector product is obtained via jax.linearize.

• The GMRES tolerance is chosen adaptively (Eisenstat–Walker style) based on the current scaled residual norm.