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]',
    rtol: 'float',
    atol: 'float | tuple[float, ...]',
    max_steps: 'int',
    dt0: 'float | 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, for example 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. 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.

Wraps the invocation boilerplate every example hand-wrote — SaveAt(ts=...), stepsize_controller(), max_steps, and — crucially — forwards self.adjoint, which most hand-written examples forgot and hardcoded diffrax.DirectAdjoint() instead. Keeping the adjoint live means SolverConfig.adjoint is honoured everywhere, and the Backsolve caveat (it cannot differentiate through values closed over in the vector field) applies as documented in ADJOINT_REGISTRY.

The crystallisation example deliberately keeps its call manual: its config sets dt0=None with a live per-trajectory step fallback, and it coerces atol to the x64 state dtype, neither of which this helper encodes. For the common case — an explicit dt0 and scalar/array tolerances — this is the whole invocation.

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; call .ys for 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  ·  also re-exported as jaxhybridmodels.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  ·  also re-exported as jaxhybridmodels.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  ·  also re-exported as jaxhybridmodels.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  ·  also re-exported as jaxhybridmodels.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