Design of callback for `scipy.integrate.solve_ivp`

Contributor & Development Discussion scipy
1

I have opened PR #21741 to add callbacks in scipy.integrate.solve_ivp for use cases such as progress tracking and dynamic max step size control. This callback is not currently intended to directly control the state or RHS of the integration as currently solve_ivp is not equipped to handle this.

The current discussion has aligned on passing the callback an instance of class IntermediateOdeResult(current naming) inheriting from _RichResult It would contain attributes that have information about the integration, e.g. time, state, nsteps, etc.

In review, there was a desire to open a discussion on when the callback should be called and what information should be supplied to the callback in those cases. I am trying to distill a few different pieces of the conversation down, and the list is meant as a starting point for discussion.

First, some definitions used in the options related to timing/type of callback:

  • after-step callback: callback that is called after every (successful) integration step of the algorithm.
  • upon-event callback: callback that is called when any event occurs
  • upon-event-type callback: callback that is called only when a specific event occurs

More definitions:

  • t_current: current integration time after step
  • y_current current integration state after step
  • t: history of time for all steps including last step
  • y: history of state for all steps including last step
  • t_events_current: times of events that occured at last step
  • y_events_current: states of system at times in above line
  • t_events history of times of events for all steps including last step
  • y_events history of states of system at times in above line
  • event_id index of a specific event that occured over last step
  • event_indices indices of all events that occured over last step
  • context what triggered the callback

Options:

  1. Only use a single after-step callback with history: include t, y, t_events, y_events, event_indices. The current values could also be passed for convenience.

  2. Only use a single after-step callback without history: include t_current, y_current, t_events_current, y_events_current, event_indices.

  3. Use a single callback, and call it after-step and upon-event without history.

    • Using context to distinguish what triggered callback.

  4. Use a single callback, and call it after-step and upon-event-type without history.

    • Similar to 3, but it would be called per step type. context could be the same as event_id in this case.

  5. Use different callbacks for after-step and upon-event without history.

    • No need for context.

  6. Use different callbacks for after-step and upon-event-type without history.

    • Need event_id and could be same as context.
    • 1 callback per event type could also be considered. event_id/context is possibly unneeded, but could be useful.

I chose to include with history for option 1, but it could potentially be considered for any option 3-6, particularly for the after-step context.

For options 3-6 we could choose to limit the information available to the different types of callback calls, e.g. only after step information in after-step context and only the state at the time of event in the upon-event context. Or we could chose to provide similar types of information to all contexts.

2

My preference/proposal is to implement option 2. And if desired, option 1 and /or option 5 or 6 could be implemented later. Option 2 is the minimal viable product for all other options except Options 3-4. Implementing Options 3-4 after Option 2 would result in the changing the number of times the callback is called.

One problem of passing the history of solution, Option 1, in the intermediate result to the callback is that it is a different shape and type than the final result in solve_ivp. The intermediate result y is a list that contains the states of the solution at each step, i.e. the shape if converted to an ndarray is (nsteps, n) where n is the number of state variables being integrated (n=len(y0)). The return from solve_ivp is an ndarray that has shape (n, nsteps). This will make it confusing to users. It could be corrected by transposing the result after every step, however this might have performance/memory issues. Additionally, if users need the history, they can construct it themselves with the callback.

For Options 3-6, they are only for convenience as all functionality can be achieved in the after-step callback context if the event information is passed. In solve_ivp the events are found after the step is completed, and the state of the solver will be the same when the callback is called in both after-step and upon-event(-type) contexts. Using Options 3-6 adds extra overhead to maintain and also describe in the documentation.

Using a single callback that is called after-step and again upon-event(-type), Options 3 and 4, adds the most complexity to the usage IMO. A single call back that is called multiple times with contexts adds more complexity to the user code if they want to do something different for after-step vs. upon-event compared to Options 5-6.

I like the idea of Options 5 and/or 6 as it reduces the complexity of the users code if they want to different things during after-step vs. upon-event(-type). However, they also increase the # of parameter in solve_ivp and thus the complexity of documentation. Since these can be done after Option 2 (or 1), it is best to wait until feedback is gained from the community before implementing it.

3

Hi @MatthewFlamm,

I’m late for this party, but here’s a minimal approach, a 10-line patch to rk.py.

callback : callable, or default None.
Called after each Runge-Kutta step with one argument, the locals() of _step_impl;
this is a dict containing t y h …, see the example below.
If the callback returns a string like “stop from callback because …”,
then solve_ivp stops, and returns sol.t sol.y … up to that point.

Pros of this approach: minimal, clear.
Cons: it’s left to the user to print / plot, log, save …

A wiki where people can show us how they use solve_ivp callbacks would be useful.
And a wiki would offload quite a few Q+As, e.g. SO scipy solve_ivp .

Notes

This patch is only for RK, not Radau etc., because tracking / displaying nd Jacobians is for experts (not me); and, I have no good test cases.

A Wibni: KeyboardInterrupt → stop solve_ivp, return sol as usual ?

Example

""" test callback in $ivp/rk.py  #bz 8jul
"""
import numpy as np
from scipy.integrate import solve_ivp

np.set_printoptions( edgeitems=5, threshold=10, linewidth=120, formatter=dict(
        float = lambda x: "%.3g" % x ))
print( 80 * "▄" )


def func( t, y ):
    return np.exp( np.exp( y ))  # fast-growing, h -> 0

func.y0 = [1]

#...............................................................................
def examplecallback( _locals ) -> None or "error msg":
    """ called with locals() from patched scipy/integrate/_ivp/rk.py """
    t = _locals["t"]
    y = _locals["y"]
    h = _locals["h"]
    hnew = _locals["h_abs"]
    stepok = _locals["step_accepted"]  # too big or too small: try again
    # rkself = _locals["self"]
    # nfev = rkself.nfev

    print( "t %-8.5g  %d  h %-6.3g  %-6.3g  y %s " % (
            t, stepok, h, hnew, y ))  # y with np.printoptions
    if hnew < 1e-8:
        return "stop from RK callback: hnew %.3g is too small  at t %g" % (hnew, t)
    maxnorm = np.linalg.norm( y, ord=np.inf )
    if maxnorm > 1e10  or np.isnan( maxnorm ):
        return "stop from RK callback: |y| %.3g is too big  at t %g" % (maxnorm, t)

#...............................................................................
sol = solve_ivp( func, [0, 1], y0=func.y0, callback=examplecallback, atol=1e-6, rtol=1e-3 )

print( f"{sol.message = } " )
print( "sol.t:", sol.t )
print( "Δt:", np.diff( sol.t ))
print( "sol.y: \n", sol.y )

Patch rk.py

166a
#...............................................................................
            if self.callback is not None:  #bz
                _ = self.callback( locals() )  # with t y h ...  #bz
                if _ is not None:  #bz
                    errstr = (_ if isinstance( _, str )  #bz
                            else "stopped by callback from RungeKutta")  #bz
                    return False, errstr  # -> sol up to now  #bz

.
103a
        self.callback = callback  #bz
.
87c
                 first_step=None, callback=None, **extraneous):  #bz
.
0a
# rk.py: class RungeKutta + callback 9jul  #bz

What do you think ?
cheers
– denis

4

There is a PR in the scipy repo related to this proposal. See link in original post. It works more or less as you have it, but with a clearer expectation on data passed to the callback IMO.

The main point of discussion in the PR and currentunresolved issue was around when the callback should be called. Assuming by your suggested patch that you also agree the callback should be called once per step, and not additionally after events (in addition to the step itself).