Stop ODE integration when a condition is satisfied?

Question

I'm simulating double pendulum fliptimes using scipy.integrate.odeint. The way my code is structured, odeint solves the system over a fixed time interval; then I have to check the result to see if a flip occurred. Instead, I want to check after each step whether the condition is satisfied, and if so, stop. Then, there is no need to solve the rest of time interval.

Here is my current code:

from matplotlib.colors import LogNorm, Normalize
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from tqdm import tqdm
import seaborn as sns
import numpy as np

def ode(y, t, length_1, length_2, mass_1, mass_2, gravity):
    angle_1, angle_1_d, angle_2, angle_2_d = y
    angle_1_dd = (-gravity * (2 * mass_1 + mass_2) * np.sin(angle_1) - mass_2 * gravity * np.sin(angle_1 - 2 * angle_2) - 2 * np.sin(angle_1 - angle_2) * mass_2 * (angle_2_d ** 2 * length_2 + angle_1_d ** 2 * length_1 * np.cos(angle_1 - angle_2))) / (length_1 * (2 * mass_1 + mass_2 - mass_2 * np.cos(2 * angle_1 - 2 * angle_2)))
    angle_2_dd = (2 * np.sin(angle_1 - angle_2) * (angle_1_d ** 2 * length_1 * (mass_1 + mass_2) + gravity * (mass_1 + mass_2) * np.cos(angle_1) + angle_2_d ** 2 * length_2 * mass_2 * np.cos(angle_1 - angle_2))) / (length_2 * (2 * mass_1 + mass_2 - mass_2 * np.cos(2 * angle_1 - 2 * angle_2)))

    return [angle_1_d, angle_1_dd, angle_2_d, angle_2_dd]

def double_pendulum(length_1, length_2, mass_1, mass_2, angle_1_init, angle_2_init, angle_1_d_init, angle_2_d_init, gravity, dt, num_steps):
    time_span = np.linspace(0, dt*num_steps, num_steps)
    y0 = [np.deg2rad(angle_1_init), np.deg2rad(angle_1_d_init), np.deg2rad(angle_2_init), np.deg2rad(angle_2_d_init)]
    sol = odeint(ode, y0, time_span, args=(length_1, length_2, mass_1, mass_2, gravity))

    return sol 

def flip(length_1, length_2, mass_1, mass_2, angle_1_init, angle_2_init, angle_1_d_init, angle_2_d_init, gravity, dt, num_steps):
    solution = double_pendulum(length_1, length_2, mass_1, mass_2, angle_1_init, angle_2_init, angle_1_d_init, angle_2_d_init, gravity, dt, num_steps)
    #angle_1 = solution[:, 0]
    angle_2 = solution[:, 2]
    for index, angle2 in enumerate(angle_2):
        if abs(angle2 - angle_2_init) >= 2*np.pi:
            return index*dt
    return dt*num_steps

angle_1_range = np.arange(-172, 172, 1)
angle_2_range = np.arange(-172, 172, 1)

fliptime_matrix = np.zeros((len(angle_1_range), len(angle_2_range)))

for i, angle_1 in tqdm(enumerate(angle_1_range), desc='angle_1'):
    for j, angle_2 in tqdm(enumerate(angle_2_range), desc='angle_2', leave=False):
        fliptime = flip(1, 1, 1, 1, angle_1, angle_2, 0, 0, 9.81, 0.01, 10000)
        fliptime_matrix[i, j] = fliptime

sns.heatmap(fliptime_matrix, square=True, cbar_kws={'label': 'Divergence'}, norm=LogNorm())
plt.xlabel('Angle 2 (degrees)')
plt.ylabel('Angle 1 (degrees)')
plt.title('Fliptime Heatmap')
plt.gca().invert_yaxis()
plt.show()

How do I stop the integration as soon as the flip condition (change in angle2 is greater than 2*pi) is satisfied?

Answer 1

In the language of initial value problems, you want to detect an "event". scipy.integrate.odeint does not provide an interface for that. This is one of the reasons the documentation suggests:

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

Once you convert your code to use solve_ivp, you can use the events parameter to terminate integration once an event is detected. The event function would look something like:

def event(t, y):
    angle2 = y[2]
    return abs(angle2 - angle_2_init) - 2*np.pi
event.terminal = True