Solve_ivp for Differential Equation

I wanted to suggest the use of scipy’s solve_ivp for the DifferentialEuqation class. Currently, odeint is used, but this is discouraged by scipy (scipy.integrate.odeint — SciPy v1.11.4 Manual):

Note

For new code, use scipy.integrate.solve_ivp to solve a differential equation.

More concrete, something like this (in class DifferentialEquation in ode/ode.py):

    def _simulate(self, y0, theta):
        # Initial condition comprised of state initial conditions and raveled sensitivity matrix
        s0 = np.concatenate([y0, self._sens_ic])
        
        # perform the integration
        sol = scipy.integrate.solve_ivp(
            fun = lambda t, Y: self._system(Y, t, tuple(np.concatenate([y0, theta]))),
            t_span=[self._augmented_times.min(), self._augmented_times.max()],
            y0=s0,
            t_eval=self._augmented_times[1:]).y.T.astype(floatX)
        
        # The solution
        y = sol[:, :self.n_states]

        # The sensitivities, reshaped to be a sequence of matrices
        sens = sol[0:, self.n_states:].reshape(self.n_times, self.n_states, self.n_p)

        return y, sens

I could open a PR on GitHub if wanted. Also, I think it would be useful that the DifferentialEquation class also allows to configure the solver. Maybe, in __init__ we could allow for some solver arguments as well?

If I recall correctly, @michaelosthege is writing some solvers to help speed up the computation. The reason I chose odeint is because we could pass different parameters to the solver at the time of solving (see the args parameter in odeint). At the time, solve_ivp only had the option to pass args to the ODE in the most recent version of scipy. We decided to stick with odeint so as not to limit dependencies.

I’m not saying this was a good decision, but it was the one we made at the time.

So far as a PR goes, I will let Mike talk more about that.