IController is the classic integral controller: dt_next = dt_used * clip(safety * E**(-1/order), factormin, factormax). For harder ODEs (mild stiffness, oscillatory error), a proportional-integral controller damps step-size oscillation by also looking at the previous step's error ratio. diffrax's preset for harder problems is PIDController(pcoeff=0.4, icoeff=0.3).
Implementation sketch
The PI update is
factor = safety * E_n**(-icoeff/order) * E_prev**(pcoeff/order)
dt_next = clip(dt_used * clip(factor, factormin, factormax), dtmin, dtmax)
with E_prev = 1.0 on the first step and E_prev updated only on accepted steps.
This is a carry-structure change, not an API change: the controller needs one scalar of state (E_prev) threaded through the bounded scan in ode.py. Cleanest shape:
- give controllers an
init(x0) -> state (returns () for ConstantStepSize/IController, E_prev0 for PIController)
- extend the contract to
adapt(x0, x1, err, dt_used, dt_prev, order, state) -> (accept, dt_next, state)
ode.py adds controller_state to the scan carry and passes it through; existing controllers ignore it, so compiled code for them is unchanged.
Everything stays inside the existing stop_gradient wrapper — the new state is controller-internal and non-differentiable by the same argument as the current comment block in controllers.py.
IController remains the default and is unchanged; PIController(rtol, atol, pcoeff=0.4, icoeff=0.3, ...) shares the clip constants and forced-accept-at-dtmin behavior.
Tests to add
PIController(pcoeff=0, icoeff=1) reproduces IController exactly (bit-for-bit rows)- oscillatory problem where PI takes fewer rejections than I at equal tolerance
- recompile:
pcoeff/icoeff as data leaves, cache size 1
- AD suite re-run with
PIController (stop-gradient still NaN-free on the flat field)
IControlleris the classic integral controller:dt_next = dt_used * clip(safety * E**(-1/order), factormin, factormax). For harder ODEs (mild stiffness, oscillatory error), a proportional-integral controller damps step-size oscillation by also looking at the previous step's error ratio. diffrax's preset for harder problems isPIDController(pcoeff=0.4, icoeff=0.3).Implementation sketch
The PI update is
with
E_prev = 1.0on the first step andE_prevupdated only on accepted steps.This is a carry-structure change, not an API change: the controller needs one scalar of state (
E_prev) threaded through the bounded scan inode.py. Cleanest shape:init(x0) -> state(returns()forConstantStepSize/IController,E_prev0forPIController)adapt(x0, x1, err, dt_used, dt_prev, order, state) -> (accept, dt_next, state)ode.pyaddscontroller_stateto the scan carry and passes it through; existing controllers ignore it, so compiled code for them is unchanged.Everything stays inside the existing
stop_gradientwrapper — the new state is controller-internal and non-differentiable by the same argument as the current comment block incontrollers.py.IControllerremains the default and is unchanged;PIController(rtol, atol, pcoeff=0.4, icoeff=0.3, ...)shares the clip constants and forced-accept-at-dtmin behavior.Tests to add
PIController(pcoeff=0, icoeff=1)reproducesIControllerexactly (bit-for-bit rows)pcoeff/icoeffas data leaves, cache size 1PIController(stop-gradient still NaN-free on the flat field)