Python scientific: interrupt differential equation solving with a condition

Question

I am currently solving a system of differential equation under python using odeint to simulate charged particles in a field (the source comes from this package ):

time = np.linspace(0, 5, 1000)

def sm(x, t): 
    return np.array([x[1], eta*Ez0(x[0])])

traj = odeint(sm,[0,1.], time)

It works fine but I would like to stop the calculation as soon as x[0] < 0. For the moment I just block the evolution of the sytem:

def sm1(x, t): 
    if x[0] < 0:
        return np.array([0, 0]) 
    else:
        return np.array([x[1], eta*Ez0(x[0])])

traj = odeint(sm1,[0,1.],time)

but I gess there are better solutions. I've found this but is seems to me that it fixes the number of steps, which is regrettable. Any suggestion appreciated.

Answer 1

If you write a custom extension of the odeint function, you could have your function raise a particular exception when it's finished. Doing it in Python might make it substantially slower, but I think you write the same thing in C or Cython. Note that I haven't tested the following.

class ThatsEnoughOfThat(Exception):
    pass

def custom_odeint(func, y0, t): # + whatever parameters you need
    for timestep in t:
        try:
            # Do stuff. Call odeint/other scipy functions?
        except ThatsEnoughOfThat:
            break
    return completedstuff

def sm2(x, t):
    if x[0] < 0:
       raise ThatsEnoughOfThat
    return np.array([x[1], eta*Ez0(x[0])])