基于GPU加速Python求解4D ODE系统的物理项目优化求助

我正在进行物理学科期末项目，需要数值求解某4D ODE系统，通过可视化初始条件随时间的演化过程展示混沌特性，目标是实现参考视频（106秒处）中的效果。但当前代码在100×100网格下运行耗时极长，且无法降低分辨率（需要展示系统的分形行为）。我使用RTX 3070显卡，花费数小时查找Python中利用GPU求解ODE的方法却几乎无果，特此寻求可行实现方案。

现有代码 #

#%% Import Packages
import numpy as np
from numpy import cos , sin
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.colors import to_rgb
import math
#%% Define Simulation parameters
TIME_FINAL = 10
FRAME_NUM = 300
MAP_RES = 1000
#%% Define Control Parameters
g = 9.8
l = 1
m = 1
#%% Define the ODE System 

# Define a time axis (number of frames)
t = np.linspace(0, TIME_FINAL, FRAME_NUM)

# Define our model (Contains the equations of motion for our system)
def model(cords , t):
    th1 , th2 , p1 , p2 = cords
    th1_dot = 6/(m*l**2) * ((2*p1 - 3*cos(th1-th2)*p2) / (16 - 9*(cos(th1 -th2))**2))
    th2_dot = 6/(m*l**2) * ((8*p2 - 3*cos(th1-th2)*p1) / (16 - 9*(cos(th1 -th2))**2))
    p1_dot = -0.5*m*l**2 * (th1_dot*th2_dot*sin(th1-th2) + 3*(g/l)*sin(th1))
    p2_dot = -0.5*m*l**2 * (-th1_dot*th2_dot*sin(th1-th2) + (g/l)*sin(th1))
    return [th1_dot , th2_dot , p1_dot , p2_dot]
    
#%% First Impression

init0 = np.array([1.3 , 0.7 , 0 , 0]) # Adjust the values to observe chaotic behivior
th1 , th2 , p1 , p2 = odeint(model,init0 ,t).T

plt.figure(dpi = 300)
plt.plot(t, th1 , label = r"\$\theta1\$", linewidth = 1)
plt.plot(t, th2 , label = r"\$\theta2\$", linewidth = 1)
plt.grid("on")
plt.legend()
plt.title(fr"\$\theta1\$ and \$\theta2\$ ({init0})")   
    
#%% Constructing an initial sub phase-space

# Define the ranges for theta and p_theta coordinates
theta1_range = np.linspace(0, 2*np.pi, MAP_RES)  
theta2_range = np.linspace(0, 2*np.pi, MAP_RES)
# Create a mesh grid of theta1 and theta2 values
theta1_grid, theta2_grid = np.meshgrid(theta1_range, theta2_range)

# %% Calculating the Sub Phase Space Progression Over Time
p1_i = 0
p2_i = 0
i = 0
phase_space = []
for (theta , phi) in zip(np.ravel(theta1_grid), np.ravel(theta2_grid)) :
    initial_condition = [theta , phi ,p1_i , p2_i]
    th1 , th2 , _ , __ = odeint(model,initial_condition ,t).T
    phase_space.append((th1,th2))
    i = i + 1
    print(f"Completed {i}/{theta1_grid.size}")

#%% Generating the Color map from the sub phase space data

def rad2rgb(theta, phi, radius = 127):
            return [math.floor(radius + radius * np.cos(phi) * np.sin(theta)),
                    math.floor(radius + radius * np.sin(phi) * np.sin(theta)),
                    math.floor(radius + radius * np.cos(theta))]
           
color_maps = []
for i in range(t.size) :
 rgb_colors = np.array([rad2rgb(phase[0][i], phase[1][i]) for phase in phase_space])
 rgb_colors = rgb_colors.reshape(theta1_grid.shape[0], theta1_grid.shape[1], 3)
 color_maps.append(rgb_colors)   
#%% Animating the Phase Space and saving into .GIF file
fig, ax = plt.subplots()

def update(frame):
    ax.clear()
    ax.imshow(color_maps[frame], origin='lower', extent=(0, 2*np.pi, 0, 2*np.pi), aspect='auto')
    ax.set_xlabel(r'\$\theta1\$')
    ax.set_ylabel(r'\$\theta2\$')
    ax.set_title(r'Color Map in \$(\theta1 , \theta2)\$ Space')


anim = FuncAnimation(fig, update, frames= len(color_maps), interval=50)
#anim.save('Fractal.gif', writer='pillow', fps=60, dpi=150)
plt.show()

针对RTX3070显卡，以下三个方案能有效利用GPU加速ODE求解，解决大规模网格的耗时问题：

方案1：JAX批量GPU加速（推荐） #

JAX支持自动向量化和GPU运算，其ODE求解器可以一次性处理所有初始条件，完全替代循环，速度提升几十倍。

步骤： #

1. 安装依赖：

pip install jax jaxlib==0.4.19+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

（根据CUDA版本调整jaxlib版本，RTX3070适配CUDA11.x）

2. 修改核心代码：

#%% 替换导入部分
import jax
import jax.numpy as jnp
from jax.experimental.ode import odeint

#%% 重定义模型为JAX兼容的向量化形式
def model(cords, t):
    th1, th2, p1, p2 = cords
    delta_th = th1 - th2
    cos_delta = jnp.cos(delta_th)
    denom = 16 - 9 * cos_delta**2
    th1_dot = 6/(m*l**2) * (2*p1 - 3*cos_delta*p2) / denom
    th2_dot = 6/(m*l**2) * (8*p2 - 3*cos_delta*p1) / denom
    p1_dot = -0.5*m*l**2 * (th1_dot*th2_dot*jnp.sin(delta_th) + 3*(g/l)*jnp.sin(th1))
    p2_dot = -0.5*m*l**2 * (-th1_dot*th2_dot*jnp.sin(delta_th) + (g/l)*jnp.sin(th1))
    return jnp.array([th1_dot, th2_dot, p1_dot, p2_dot])

#%% 批量处理所有初始条件
# 构造形状为(MAP_RES*MAP_RES, 4)的初始条件数组
initial_conditions = jnp.stack([
    jnp.ravel(theta1_grid),
    jnp.ravel(theta2_grid),
    jnp.full(MAP_RES*MAP_RES, p1_i),
    jnp.full(MAP_RES*MAP_RES, p2_i)
], axis=1)

# 一次性求解所有初始条件，结果形状为(FRAME_NUM, MAP_RES*MAP_RES, 4)
solutions = odeint(model, initial_conditions, t)

# 提取th1和th2，形状为(FRAME_NUM, MAP_RES*MAP_RES)
th1_all = solutions[:, :, 0]
th2_all = solutions[:, :, 1]

#%% 向量化生成颜色映射（GPU加速）
def rad2rgb_jax(theta, phi, radius=127):
    r = radius + radius * jnp.cos(phi) * jnp.sin(theta)
    g = radius + radius * jnp.sin(phi) * jnp.sin(theta)
    b = radius + radius * jnp.cos(theta)
    return jnp.stack([r, g, b], axis=-1).astype(jnp.int32)

# 生成所有帧的颜色图，形状为(FRAME_NUM, MAP_RES, MAP_RES, 3)
color_maps = jnp.array([
    rad2rgb_jax(th1_all[i], th2_all[i]).reshape(MAP_RES, MAP_RES, 3)
    for i in range(FRAME_NUM)
])

# 转成NumPy数组用于Matplotlib动画
color_maps_np = np.array(color_maps)

方案2：CuPy+SciPy混合加速 #

CuPy将NumPy运算转移到GPU，虽然SciPy的odeint是CPU的，但可以把模型内的计算放在GPU上，同时通过批量初始条件减少循环开销。

步骤： #

1. 安装依赖：

pip install cupy-cuda11x

2. 修改核心代码：

#%% 替换导入
import cupy as cp
from scipy.integrate import odeint

#%% 模型函数使用CuPy
def model_cupy(cords, t):
    th1, th2, p1, p2 = cp.array(cords)
    delta_th = th1 - th2
    cos_delta = cp.cos(delta_th)
    denom = 16 - 9 * cos_delta**2
    th1_dot = 6/(m*l**2) * (2*p1 - 3*cos_delta*p2) / denom
    th2_dot = 6/(m*l**2) * (8*p2 - 3*cos_delta*p1) / denom
    p1_dot = -0.5*m*l**2 * (th1_dot*th2_dot*cp.sin(delta_th) + 3*(g/l)*cp.sin(th1))
    p2_dot = -0.5*m*l**2 * (-th1_dot*th2_dot*cp.sin(delta_th) + (g/l)*cp.sin(th1))
    return cp.asnumpy(cp.array([th1_dot, th2_dot, p1_dot, p2_dot]))

#%% 批量准备初始条件（CuPy数组）
initial_conditions_cp = cp.stack([
    cp.ravel(theta1_grid),
    cp.ravel(theta2_grid),
    cp.full(MAP_RES*MAP_RES, p1_i),
    cp.full(MAP_RES*MAP_RES, p2_i)
], axis=1)

# 循环求解（比原NumPy循环快很多，因为模型计算在GPU）
phase_space = []
for ic in initial_conditions_cp:
    sol = odeint(model_cupy, cp.asnumpy(ic), t)
    phase_space.append((sol[:,0], sol[:,1]))

方案3：PyTorch+torchdiffeq #

torchdiffeq是PyTorch生态的ODE求解库，原生支持GPU批量求解，适合熟悉PyTorch的用户。

步骤： #

1. 安装依赖：

pip install torch torchdiffeq

2. 修改核心代码：

#%% 替换导入
import torch
from torchdiffeq import odeint as torch_odeint

#%% 模型函数转为PyTorch兼容形式
class ODEModel(torch.nn.Module):
    def forward(self, t, cords):
        th1, th2, p1, p2 = cords.unbind(dim=-1)
        delta_th = th1 - th2
        cos_delta = torch.cos(delta_th)
        denom = 16 - 9 * cos_delta**2
        th1_dot = 6/(m*l**2) * (2*p1 - 3*cos_delta*p2) / denom
        th2_dot = 6/(m*l**2) * (8*p2 - 3*cos_delta*p1) / denom
        p1_dot = -0.5*m*l**2 * (th1_dot*th2_dot*torch.sin(delta_th) + 3*(g/l)*torch.sin(th1))
        p2_dot = -0.5*m*l**2 * (-th1_dot*th2_dot*torch.sin(delta_th) + (g/l)*torch.sin(th1))
        return torch.stack([th1_dot, th2_dot, p1_dot, p2_dot], dim=-1)

#%% 批量求解
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ODEModel().to(device)

# 构造初始条件张量，形状(MAP_RES*MAP_RES, 4)
initial_conditions = torch.tensor(np.stack([
    np.ravel(theta1_grid),
    np.ravel(theta2_grid),
    np.full(MAP_RES*MAP_RES, p1_i),
    np.full(MAP_RES*MAP_RES, p2_i)
], axis=1), dtype=torch.float32).to(device)

t_tensor = torch.tensor(t, dtype=torch.float32).to(device)

# 求解，结果形状(FRAME_NUM, MAP_RES*MAP_RES, 4)
solutions = torch_odeint(model, initial_conditions, t_tensor, method='rk4')

# 提取th1和th2
th1_all = solutions[:, :, 0].cpu().numpy()
th2_all = solutions[:, :, 1].cpu().numpy()

• 颜色映射部分完全用向量化操作替代循环，避免Python overhead
• 动画生成时，提前将所有帧数据加载到内存，避免实时计算
• 如果MAP_RES=1000，建议分批次求解（比如每次处理1000个初始条件），避免GPU内存溢出

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