加速Python微分方程求解运算速度——摆-车系统模拟优化

我需要求解自由移动小车上摆锤的位置随时间变化的关系，推导了运动方程并用Python模拟验证。程序能正常运行并绘制曲线，但提升绘图精度时速度极慢，作为Python和Scipy新手，求优化建议。

原程序代码 #

from scipy.integrate import quad
from scipy.optimize import fsolve
import numpy as np
import matplotlib.pyplot as plt

# These values can be changed
masstot = 5
mass = 2
g= 9.8
l = 9.8
wan = (g/l)**(1/2)
vuk = 0.1
oug = 1

def afad(lah): # Find first constant
    wan = 1
    vuk = 0.1
    oug = 1
    kan = (12*(lah**4)*((3*(vuk**2)-(wan**2))))-((16*((wan**2)-(vuk**2))-(5*oug**2))*(lah**2))+(4*(oug**2))
    return (kan)

solua = fsolve(afad, 1)

intsolua = sum(solua) 

def kfad(solua, wan, vuk): # Find second constant
    res = ((wan**2)-(vuk**2)-((2*(solua**2)*((2*(vuk**2))+(wan**2)))/((5*(solua**2))+4)))**(1/2)
    return (res)

ksol = kfad(solua, wan, vuk)

def deg(t, solua, vuk, ksol): # Find angle of pendulum relative to time
    res = 2*np.arctan(solua*np.exp(-1*vuk*t)*np.sin(ksol*t))
    return(res)

def chandeg(t, solua, vuk, ksol): # Find velocity of pendulum relative to time
    res = (((-2*solua*vuk*np.exp(vuk*t)*np.sin(ksol*t))+(2*solua*ksol*np.exp(vuk*t)*np.cos(ksol*t)))/(np.exp(2*vuk*t)+((solua**2)*(np.sin(ksol*t)**2))))
    return(res)

xs = np.linspace(0, 60, 20) # Value can be changed to alter plotting accuracy  and length

def dinte1(deg, bond, solua, vuk, ksol): # used to plot angle at at a certain time
    res = []
    for x in (bond):
        res.append(deg(x, solua, vuk, ksol))
    return res

def dinte2(chandeg, bond, solua, vuk, ksol): # used to plot angular velocity at a certain time
    res = []
    for x in (bond):
        res.append(chandeg(x, solua, vuk, ksol))
    return res

def dinte(a, bond, mass, l, solua, vuk, ksol, g, masstot ): # used to plot acceleration of system at certain time
    res = []
    for x in (bond):
        res.append(a(x, mass, l, solua, vuk, ksol, g, masstot))
    return res

def a(t, mass, l, solua, vuk, ksol, g, masstot): # define acceleration of system to time
    return (((mass*l*(chandeg(t, solua, vuk, ksol)**2))+(mass*g*np.cos(deg(t, solua, vuk, ksol))))*np.sin(deg(t, solua, vuk, ksol))/masstot)

def j(t):
    return sum(a(t, mass, l, intsolua, vuk, ksol, g, masstot))

def f(ub):
    return quad(lambda ub: quad(j, 0, ub)[0], 0, ub)[0]

def int2(f, bond): # Integrates system acceleration twice to get posistion relative to time
    res = []
    for x in (bond):
        res.append(f(x))
        print(res)
    return res

plt.plot(xs, int2(f, xs)) # This part of the program runs quite slowly
#plt.plot(xs, dinte(a, xs, mass, l, solua, vuk, ksol, g, masstot))
#plt.plot(xs, dinte2(chandeg, xs, solua, vuk, ksol))
#plt.plot(xs, dinte1(deg, xs, solua, vuk, ksol))
plt.show()

核心优化建议 #

1. 用NumPy向量运算替代循环 #

原代码中dinte1、dinte2、dinte都是通过for循环逐个计算数组元素，完全可以利用NumPy的广播机制直接对整个数组进行运算，大幅提升速度。比如deg函数可以直接接收数组t，无需循环：

# 直接处理数组t，无需循环
theta = deg(xs, solua[0], vuk, ksol[0])

2. 避免重复计算冗余值 #

在a函数中，多次调用deg(t,...)和chandeg(t,...)，可以先计算一次这两个值再复用，减少计算开销：

def a(t, mass, l, solua, vuk, ksol, g, masstot):
    theta = deg(t, solua, vuk, ksol)
    theta_dot = chandeg(t, solua, vuk, ksol)
    return ((mass*l*(theta_dot**2) + mass*g*np.cos(theta)) * np.sin(theta)) / masstot

3. 替换低效的二重积分方式 #

原代码中f(ub)用嵌套quad逐个计算每个时间点的积分，这是速度慢的核心原因。建议用scipy.integrate.cumulative_trapezoid一次性计算整个时间序列的累积积分：

from scipy.integrate import cumulative_trapezoid

# 先计算整个时间序列的加速度
acc = a(xs, mass, l, solua[0], vuk, ksol[0], g, masstot)
# 积分一次得到速度（初始速度设为0）
vel = cumulative_trapezoid(acc, xs, initial=0)
# 再积分一次得到位置（初始位置设为0）
pos = cumulative_trapezoid(vel, xs, initial=0)

4. 移除不必要的操作 #

• 删除j函数中的sum：a函数返回单个值，sum无意义，直接返回a(...)即可
• 删除int2函数中的print(res)：打印操作会大幅拖慢大数组的处理速度
• 修正afad函数：函数内重新定义的wan、vuk、oug覆盖了全局参数，导致计算错误，应删除这些局部变量，使用全局定义的值

5. 优化常数处理 #

fsolve返回的是单元素数组，计算ksol时直接用solua[0]代替整个数组，避免不必要的数组运算：

ksol = kfad(solua[0], wan, vuk)

优化后的完整代码 #

from scipy.optimize import fsolve
from scipy.integrate import cumulative_trapezoid
import numpy as np
import matplotlib.pyplot as plt

# 可配置参数
masstot = 5
mass = 2
g = 9.8
l = 9.8
wan = np.sqrt(g / l)
vuk = 0.1
oug = 1

def afad(lah):
    # 使用全局参数，不再重复定义
    kan = (12 * (lah**4) * (3 * vuk**2 - wan**2)) - \
          ((16 * (wan**2 - vuk**2) - 5 * oug**2) * lah**2) + \
          4 * oug**2
    return kan

solua = fsolve(afad, 1)[0]  # 直接取单元素值

def kfad(solua, wan, vuk):
    res = np.sqrt(wan**2 - vuk**2 - (2 * solua**2 * (2 * vuk**2 + wan**2)) / (5 * solua**2 + 4))
    return res

ksol = kfad(solua, wan, vuk)

def deg(t, solua, vuk, ksol):
    return 2 * np.arctan(solua * np.exp(-vuk * t) * np.sin(ksol * t))

def chandeg(t, solua, vuk, ksol):
    numerator = -2 * solua * vuk * np.exp(vuk * t) * np.sin(ksol * t) + 2 * solua * ksol * np.exp(vuk * t) * np.cos(ksol * t)
    denominator = np.exp(2 * vuk * t) + solua**2 * np.sin(ksol * t)**2
    return numerator / denominator

def a(t, mass, l, solua, vuk, ksol, g, masstot):
    theta = deg(t, solua, vuk, ksol)
    theta_dot = chandeg(t, solua, vuk, ksol)
    return ((mass * l * theta_dot**2 + mass * g * np.cos(theta)) * np.sin(theta)) / masstot

# 生成时间序列，可调整点数提升精度
xs = np.linspace(0, 60, 1000)  # 这里设1000点也能快速运行

# 计算加速度、速度、位置
acc = a(xs, mass, l, solua, vuk, ksol, g, masstot)
vel = cumulative_trapezoid(acc, xs, initial=0)
pos = cumulative_trapezoid(vel, xs, initial=0)

# 绘图
plt.plot(xs, pos)
# plt.plot(xs, acc)  # 可按需开启其他曲线
# plt.plot(xs, deg(xs, solua, vuk, ksol))
plt.show()

内容的提问来源于stack exchange，提问作者Juan A