在满足条件时停止ODE积分？

Question

我正在用 scipy.integrate.odeint 模拟双摆的翻转时间。我的代码结构是这样的：odeint 会在一个固定的时间段内解决这个系统；然后我得检查结果，看是否发生了翻转。不过，我想要的是在每一步之后检查一下条件是否满足，如果满足就停止。这样就不需要继续计算剩下的时间段了。

这是我现在的代码：

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()

我该如何在翻转条件（angle2 的变化大于 2*pi）满足时立即停止计算呢？

Answer 1

在处理初值问题时，你可能想要检测某个“事件”。但是，scipy.integrate.odeint 这个工具并没有提供检测事件的功能。这也是文档建议你使用其他方法的原因之一：

对于新代码，建议使用 scipy.integrate.solve_ivp 来解决微分方程。

一旦你把代码转换为使用 solve_ivp，你就可以利用 events 这个参数，在检测到事件后停止计算。事件函数大致会是这样的：

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