Solver: ODE Integration

SolverConfig bundles a diffrax solver instance with its tolerances and step controls. Every field is static, so the config is closed over by jitted functions without re-tracing on value changes (a tolerance change does trigger a recompile, which is what we want).

Solvers are looked up by name through SOLVER_REGISTRY; register_solver extends the registry with custom implementations so saved configs round-trip cleanly.

Quick links

• SolverConfig
• SOLVER_REGISTRY
• register_solver
• ADJOINT_REGISTRY
• register_adjoint

---

SolverConfig

from jaxhybridmodels.solver import SolverConfig  ·  also re-exported as jaxhybridmodels.SolverConfig

SolverConfig(
    solver: 'diffrax.AbstractSolver[Any] | None' = None,
    rtol: 'float' = 1e-05,
    atol: 'float | tuple[float, ...]' = 1e-07,
    max_steps: 'int' = 4096,
    dt0: 'float | None' = None,
    adjoint: 'diffrax.AbstractAdjoint | None' = None,
    pcoeff: 'float' = 0.0,
    icoeff: 'float' = 1.0,
    dcoeff: 'float' = 0.0,
) -> None

Everything the ODE solve needs, held as static configuration.

Every field is eqx.field(static=True), so the config carries no JAX array leaves. Compiled functions close over it, so changing a value recompiles rather than reusing the old kernel. That is intended: a tolerance change must change the compiled solve.

Attributes

Field	Type	Description
solver	diffrax.AbstractSolver	Concrete solver instance, default diffrax.Tsit5(). Its class must appear in SOLVER_REGISTRY for to_dict to round-trip.
rtol, atol	`float	tuple[float, ...]`
max_steps	int	Upper limit on solver steps, default 4096. The solve errors rather than running forever if it needs more.
dt0	`float	None`
adjoint	diffrax.AbstractAdjoint	How gradients are taken back through the solve. See ADJOINT_REGISTRY for what each choice costs.
pcoeff, icoeff, dcoeff	float	Gains of the PID step-size controller. The defaults (0, 1, 0) are diffrax's own and give plain I-control. See stepsize_controller.

Source

SolverConfig.diffeqsolve()

diffeqsolve(
    self,
    term: 'diffrax.AbstractTerm',
    ts: 'Array',
    y0: 'Array',
    args: 'Any' = None,
) -> diffrax.Solution

Run the solve this config describes, so simulate_fn stays thin.

Calls diffrax.diffeqsolve with this config's solver, dt0, stepsize_controller(), max_steps and adjoint, saving at ts. Forwarding adjoint matters: a hand-written call that hardcodes diffrax.DirectAdjoint() silently ignores SolverConfig.adjoint. The Backsolve caveat (it cannot differentiate through values closed over in the vector field) is documented in ADJOINT_REGISTRY.

Call diffrax.diffeqsolve yourself when you need something this helper does not pass, such as a per-trajectory dt0. The crystallisation example does this.

Parameters

Parameter	Type	Description
term		The diffrax term, usually diffrax.ODETerm(vector_field) where vector_field is the user's physics. The user still owns the physics; this folds only the invocation.
ts		The observation times [T]. The solver integrates from ts[0] to ts[-1] and saves at exactly ts.
y0		Full initial state [S].
args		Optional static-or-traced value passed to the vector field's args (e.g. the predictors pytree, for Backsolve).

Returns

Item	Type	Description
diffrax.Solution		The diffrax solution; .ys is the state trajectory [T, S] simulate_fn must return.

Source

SolverConfig.stepsize_controller()

stepsize_controller(self) -> 'diffrax.PIDController'

Build the adaptive step-size controller this config describes.

Removes a coercion every caller had to remember: diffrax broadcasts atol against the state pytree, and a Python tuple is not an array, so per-state tolerances misbehaved unless the caller wrapped them in jnp.asarray first.

The (0, 1, 0) coefficient defaults are diffrax's own plain I-control, so this reproduces the PIDController(rtol, atol) the examples wrote by hand. Raise pcoeff to 0.3 or 0.4 to damp step-size oscillation on stiff problems.

Source

SolverConfig.to_dict()

to_dict(self) -> 'dict[str, Any]'

Serialise to a JSON-compatible dict via the two registries.

Solver and adjoint instances become their registered names, and a tuple atol becomes a list. An unregistered class raises rather than being guessed at.

Source

---

SOLVER_REGISTRY

from jaxhybridmodels.solver import SOLVER_REGISTRY

SOLVER_REGISTRY = {
  'Dopri5': _MetaAbstractSolver
  'Heun': _MetaAbstractSolver
  'Kvaerno3': _MetaAbstractSolver
  'Tsit5': _MetaAbstractSolver
}

dict() -> new empty dictionary

dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

---

register_solver()

from jaxhybridmodels.solver import register_solver

register_solver(name: 'str', cls: 'type[diffrax.AbstractSolver[Any]]') -> 'None'

Register a custom diffrax solver class under name for round-trip serialisation.

After registration, SolverConfig(solver=cls(), ...).to_dict() emits {"solver": name, ...} and SolverConfig.from_dict accepts it. Re-registering an existing name overwrites without warning. Calling code owns the naming.

Source

---

ADJOINT_REGISTRY

from jaxhybridmodels.solver import ADJOINT_REGISTRY

ADJOINT_REGISTRY = {
  'Backsolve': _ModuleMeta
  'Direct': _ModuleMeta
  'ForwardMode': _ModuleMeta
  'RecursiveCheckpoint': _ModuleMeta
}

dict() -> new empty dictionary

dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

---

register_adjoint()

from jaxhybridmodels.solver import register_adjoint

register_adjoint(name: 'str', cls: 'type[diffrax.AbstractAdjoint]') -> 'None'

Register a diffrax adjoint class under name for round-trip serialisation.

Same contract as register_solver. Re-registering an existing name overwrites without warning.

Source