Python大规模ODE系统求解提速、内存优化及生成器使用咨询

你编写的生成器逻辑完全正确，每次迭代仅返回当前s参数对应的稳态解，没有在生成器内部累积存储全量结果，这个设计本身已经符合节省内存的需求。

核心原因是dxidt右端函数被ODE求解器高频调用时，反复生成大体积临时数组，加上全量存储所有节点的解导致内存溢出，优化点如下：

• 移除dxidt内部的重复计算：past_index数组在m和节点数固定时完全不变，不要每次调用dxidt都重新生成，提前在循环外计算好作为参数传入即可，这是解决内存爆炸的核心改动
• 减少不必要的全量存储：如果不需要保留所有s参数下全部节点的稳态值，仅需统计分层和、最大值等聚合指标，直接在迭代生成器时计算指标存储即可，不需要把整个解数组存入eqsol2，内存占用可降低90%以上
• 提前计算固定参数：sum(m[i]**np.arange(0, n))这类每个m值对应的固定参数，提前计算好存入变量，不要每次s循环都重复生成临时数组

• 开启Numba无python模式：把@numba.jit改成@numba.jit(nopython=True)，可带来数倍到数十倍的速度提升，避免Numba退化成纯Python运行模式
• 提前计算重复值：aj每行的和np.sum(aj[i])在aj固定时不会变化，提前算好传入dxidt，避免每次调用右端函数都重复求和
• 优化求解器设置：你只需要稳态解，不需要中间时间步结果，调用odeint时仅传入首尾两个时间点即可，不需要存100步的中间结果；如果系统是刚性的，换成solve_ivp的BDF或Radau求解器，收敛速度会快很多
• 移除高频打印：s循环内的print操作非常拖慢速度，大规模计算时注释掉即可，仅保留每个m循环的一次打印做进度提示

import numba
import numpy as np
import scipy as sci
from scipy.integrate import odeint

# 优化后的右端函数，无重复临时数组生成
@numba.jit(nopython=True)
def dxidt(xi, t, ao, aj_sum, past_index, d):
    dx = np.zeros_like(xi)
    for i in range(1, len(dx)):
        dx[i] = ao[i-1] * xi[past_index[i-1]] - (d + aj_sum[i]) * xi[i]
    return dx

# 优化后的生成器
def ode_generator(d, m_arr, n, s, t):
    t_trim = t[[0, -1]] # 仅保留首尾时间点，不需要中间结果
    for m in m_arr:
        # 每个m值对应的固定参数只算一次
        node_count = np.sum(m ** np.arange(0, n))
        aj_total_len = np.sum(m ** np.arange(0, n+1)) - 1
        past_index = np.repeat(np.arange(0, node_count), m)
        xi = np.ones(node_count)
        aj = np.ones(aj_total_len)
        tau = np.random.uniform(low=0., high=1., size=aj_total_len)
        tau[:m] = 1.
        print(f"处理m={m}")
        for s_val in s:
            aj[m:] = 1 + s_val * tau[m:]
            ao = aj[:node_count - 1]
            ajt = aj.reshape(-1, m)
            aj_sum = ajt.sum(axis=1)
            sol = odeint(dxidt, xi, t_trim, args=(ao, aj_sum, past_index, d))[-1]
            # 若不需要全量sol，直接返回聚合指标即可，例如 yield sol.sum()
            yield sol

# 原参数部分修改m的变量名避免冲突
d = 1.
m_arr = np.array([2,4,8], dtype=int)
n = 6
s = np.concatenate([np.arange(0., 0.1, 0.001),
                    np.arange(0.1, 1., 0.01),
                    np.arange(1., 10., 0.1),
                    np.arange(10., 100., 1.),
                    np.arange(100., 1010., 10.)])
tf = 10
steps = 100
t = np.linspace(0, tf, steps)

# 迭代计算，若只存指标可直接在这里处理
eqsol2 = []
for sol in ode_generator(d, m_arr, n, s, t):
    eqsol2.append(sol)

内容的提问来源于stack exchange，提问作者Irbin B.