I have a general function that defines a form of an ODE that I plan to integrate using scipy.integrate.odeint, for example:

def my_ode(K, tau, y, u):
  return K*u/tau - y/tau  # dydt

I have several objects in my code that all have dynamics of the form defined in my_ode, but with unique parameters K and tau. I would love to be able to just pass a unique handle to my_ode with those parameters already set when I initialize my objects, so that when I update my objects, all I have to do is something like soln = odeint(my_ode, t, y, u) for some simulation time t.

For example, if I define a class:

class MyThing:
  def __init__(self, ode, y0):
    # I would rather not maintain K and tau in the objects, I just want the ODE with unique parameters here.
    self.ode = ode
    self.y = y0
    self.time = 0.0

  def update(self, t, u):
    # I want this to look something like:
    self.y = scipy.integrate.odeint(self.ode, t, self.y, u)

Can I do something with Lambdas when I initialize instances of MyThing to basically assign parameters K and tau at initialization and never need to pass them again? I am a bit stuck.

Solution with lambdas

It looks like I can make this work using lambdas to generate unique function handles when I initialize my objects. For compatibility with odeint, I need to define my functions so that the first two arguments are time and initial state:

def my_ode(t, y, u, K, tau):
  return K*u/tau - y/tau  # dydt

Next I can initialize objects of MyThing using lambdas to set K and tau as:

thing1 = MyThing(lambda t, y, u: my_ode(t, y, u, 10.0, 0.5), 0.0)

The function handle that gets assigned to thing1.ode is now the function handle returned by the lambda (this may not be the right way to say this) with values for K and tau set. Now in thing1.update, I need to make some changes to get it to work with odeint:

def update(self, t_step, t_end, u):
  t_array = np.arange(self.time, t_end, t_step)  # time values at which to evaluate ODE
  response = scipy.integrate.odeint(self.ode, self.y, t_array, (u,))
  self.y = response[-1]  # current state is the last evaluated state

One thing that tripped me up a bit is that any extra arguments to your ODE need to be passed as a tuple to odeint. This seems to work pretty well for what I want.

There is also the more object-oriented approach using scipy.integrate.ode, which allows for step-wise integration of the function and is great for my simulation purposes. For this, I set the object's ODE and update it with something like:

class MyThing():
  def __init__(self, ode, y0):
    self.ode = integrate.ode(ode)  # define the ODE
    self.ode.set_integrator("dopri5")  # choose an integrator
    self.ode.set_initial_value(y0)

  def update(self, u, t_step):
    """Update the ODE step-wise."""
    self.ode.set_f_params(u)  # need to pass extra parameters with this method
    self.ode.integrate(self.ode.t + t_step)  # step-wise update
    return self.ode.successful()
  
  def get_output(self):
    """Get output from ODE function."""
    return self.ode.y