Steady state solver termination

Hi Patrick,

I’ve been trying to implement the built-in steady_state_event to terminate the solver when the system reaches a steady state, but I haven’t had much luck so far. It seems that the solver doesn't terminate as expected.

Here is the code I’ve been using:

sol = diffrax.diffeqsolve(
    terms=diffrax.ODETerm(model.rhs_jit),
    solver=diffrax.Tsit5(),
    t0=0.0,
    t1=jnp.inf,
    dt0=None,
    max_steps=None,
    event=diffrax.Event(diffrax.steady_state_event),
    y0=y0,
    throw=True,
    progress_meter=diffrax.TqdmProgressMeter(),
    saveat=diffrax.SaveAt(ts=jnp.linspace(0.0, t_end, n_steps)),
    stepsize_controller=diffrax.PIDController(rtol=1e-4, atol=1e-6),
    args=(F_omega, F_amp),
)

I want to create a custom steady_state_event where I sample the system's response at intervals based on the oscillation period (which is equal to the excitation frequency). By comparing the last measured sample with the previous one (within a specified tolerance), I can determine when the system has reached steady state.

However, I’m having difficulty understanding how the event mechanism works. When I provide my own cond_fn, I get the function parameters like t, state, and args. But it seems that the time and state doesn't contain the history of the solved states at time t. This history is essential for determining steady state, as I need to compare the most recent state with previous ones.

I noticed something in the documentation about using the adjoint implicit function theorem for backpropagating through the solution, which might be helpful, but I’m unsure how to implement it in my case.

Here’s the current implementation of my custom steady_state_event function:

def _steady_state_event(t, state, args, **kwargs) -> jax.Array:
    del kwargs
    
    # My steady state check logic here
    
    return False

And the solver call:

sol = diffrax.diffeqsolve(
    terms=diffrax.ODETerm(model.rhs_jit),
    solver=diffrax.Tsit5(),
    t0=0.0,
    t1=t_end,
    dt0=None,
    max_steps=4096,
    event=diffrax.Event(_steady_state_event),
    y0=y0,
    throw=True,
    progress_meter=diffrax.TqdmProgressMeter(),
    saveat=diffrax.SaveAt(ts=jnp.linspace(0.0, t_end, n_steps)),
    stepsize_controller=diffrax.PIDController(rtol=1e-4, atol=1e-6),
    args=(F_omega, F_amp),
)

Please keep in mind that in my use case I vmap everything. And also I am sure that my solution reaches steady state.

Could you help clarify how I can properly implement this steady state check and resolve the issues I’m encountering with the event mechanism?

Thanks in advance!

Start with diffrax.Event, diffrax.steady_state_event, and the custom _steady_state_event shown in the issue, then read the event documentation and inspect how t, state, and args are supplied. Verify how a history-dependent check could work with vmap and confirm that completion is reported when the sampled states meet the requested tolerance.