NumPy odeint output extra variables

Question

What is the easiest way to save intermediate variables during simulation with odeint in Numpy?

For example:

def dy(y,t)
    x = np.rand(3,1)
    return y + x.sum()

sim = odeint(dy,0,np.arange(0,1,0.1))

What would be the easiest way to save the data stored in x during simulation? Ideally at the points specified in the t argument passed to odeint.

Answer 1

A handy way to hack odeint, with some caveats, is to wrap your call to odeint in method in a class, with dy as another method, and pass self as an argument to your dy function. For example,

class WrapODE(object):
    def __init__(self):
        self.y_0 = 0.
        self.L_x = []
        self.timestep = 0
        self.times = np.arange(0., 1., 0.1)

    def run(self):
        self.L_y = odeint(
            self.dy,
            self.y_0, self.times,
            args=(self,))

    @staticmethod
    def dy(y, t, self):
        """"
        Discretized application of dudt

        Watch out! Because this is a staticmethod, as required by odeint, self
        is the third argument
        """
        x = np.random.rand(3,1)
        if t >= self.times[self.timestep]:
            self.timestep += 1
            self.L_x.append(x)
        else:
            self.L_x[-1] = x
        return y + x.sum()

To be clear, this is a hack that is prone to pitfalls. For example, unless odeint is doing Euler stepping, dy is going to get called more times than the number of timesteps you specify. To make sure you get one x for each y, the monkey business in the if t >= self.times[self.timestep]: block picks a spot in an array for storing data for each time value from the times vector. Your particular application might lead to other crazy problems. Be sure to thoroughly validate this method for your application.