Python双摆仿真大规模计算避免内存错误的解决方案

核心问题分析 #

scipy的odeint默认以float64类型返回全时间序列的解，哪怕输入的t是float32。对于10亿+1个采样点的双摆系统（4维状态：两个摆角、两个角速度），总数据量约为1e9+1 * 4 * 8字节 ≈ 32GiB，远超普通机器内存容量，必然触发MemoryError。

可行解决方法 #

1. 分块处理时间序列，逐段计算并落地磁盘 #

无需一次性计算所有时间点，将大时间序列拆分为多个子区间，每次仅计算一段结果并写入磁盘，计算完成后立即释放内存：

import numpy as np
from scipy.integrate import odeint

# 双摆微分方程实现（示例）
def double_pendulum(y, t, l1, l2, m1, m2, g):
    theta1, theta1_dot, theta2, theta2_dot = y
    # 双摆动力学方程推导结果
    dydt = [
        theta1_dot,
        (-g*(2*m1 + m2)*np.sin(theta1) - m2*g*np.sin(theta1-2*theta2) - 2*np.sin(theta1-theta2)*m2*(theta2_dot**2*l2 + theta1_dot**2*l1*np.cos(theta1-theta2))) / (l1*(2*m1 + m2 - m2*np.cos(2*theta1-2*theta2))),
        theta2_dot,
        (2*np.sin(theta1-theta2)*(theta1_dot**2*l1*(m1+m2) + g*(m1+m2)*np.cos(theta1) + theta2_dot**2*l2*m2*np.cos(theta1-theta2))) / (l2*(2*m1 + m2 - m2*np.cos(2*theta1-2*theta2)))
    ]
    return dydt

# 初始条件与系统参数
y0 = [np.pi/4, 0, np.pi/4, 0]
params = (1.0, 1.0, 1.0, 1.0, 9.8)

total_points = 1000000001
total_duration = 10000  # 模拟总时长，按需调整
chunk_size = 1000000    # 每段采样点数量，根据内存容量调整

y_prev_last = y0
for i in range(0, total_points, chunk_size):
    # 生成当前段的时间点
    start_t = i * total_duration / total_points
    end_t = min((i + chunk_size) * total_duration / total_points, total_duration)
    t_chunk = np.linspace(start_t, end_t, min(chunk_size, total_points - i), dtype=np.float32)
    
    # 计算当前段的解
    y_chunk = odeint(double_pendulum, y_prev_last, t_chunk, args=params)
    # 保存结果到磁盘，自动释放内存
    np.savez(f"pendulum_chunk_{i//chunk_size}.npz", t=t_chunk, y=y_chunk)
    # 记录当前段最后一个状态，作为下一段初始条件
    y_prev_last = y_chunk[-1]

2. 改用solve_ivp配合回调函数实时处理数据 #

solve_ivp是scipy更现代的ODE求解器，支持回调函数，可在每个采样点实时处理数据，无需保存全量序列：

from scipy.integrate import solve_ivp
import numpy as np

# 回调函数：实时处理或保存单个采样点数据
def process_sample(t, y):
    # 将数据追加写入文件，避免内存累积
    with open("pendulum_data.txt", "ab") as f:
        np.savetxt(f, np.hstack([t, y]).reshape(1, -1))
    return False  # 返回False表示继续求解

# 配置求解参数
t_span = (0, 10000)
t_eval = np.linspace(0, 10000, 1000000001, dtype=np.float32)

# 启动求解，触发回调处理每个采样点
sol = solve_ivp(double_pendulum, t_span, y0, args=params, 
                t_eval=t_eval, method='RK45', callback=process_sample)

3. 强制使用float32降低内存占用 #

若计算精度允许，将结果转换为float32后保存，可使内存/存储需求减半：

# 在分块处理代码中，保存前转换数据类型
y_chunk_float32 = y_chunk.astype(np.float32)
np.savez(f"pendulum_chunk_{i//chunk_size}.npz", t=t_chunk, y=y_chunk_float32)

4. 在线计算概率密度，不保存原始数据 #

若目标是生成状态概率密度，无需保存所有采样点，可实时更新直方图或统计量：

# 初始化4维状态的直方图 bins
bins = [np.linspace(-np.pi, np.pi, 100), 
        np.linspace(-10, 10, 100), 
        np.linspace(-np.pi, np.pi, 100), 
        np.linspace(-10, 10, 100)]
hist, _ = np.histogramdd([], bins=bins)

# 分块更新直方图
for i in range(0, total_points, chunk_size):
    # ... 计算y_chunk ...
    current_hist, _ = np.histogramdd(y_chunk, bins=bins)
    hist += current_hist

# 归一化得到概率密度
prob_density = hist / total_points

额外建议 #

• 双摆是混沌系统，若无需极高时间分辨率，可适当减少采样点数量（比如从10亿降至1亿），内存需求会降至3GiB左右，普通机器即可处理。
• 超大规模采样可考虑分布式计算框架（如Dask）拆分任务，或使用GPU加速求解（如CuPy配合GPU版ODE求解器）。

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