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?