使用SciPy odeint求解ODE遇NaN及积分超限问题求助

我来帮你拆解下遇到的问题，然后给出具体的解决办法：

问题根源分析 #

你遇到的两个核心问题，本质是数值求解的特性和函数鲁棒性不足导致的：

• odeint自适应步长越界：虽然你定义的mlist2是从x1*m0到m0，但scipy的odeint默认用自适应步长迭代，过程中可能会试探性让m超过m0。而你的sigma_sq是单调递减函数，m>m0时sigma_sq(m)-sigma_sq(m0)会变成负数，自然无法取3/2次幂。
• 数值稳定性与函数鲁棒性问题：当m接近m0时，sigma_sq(m)-sigma_sq(m0)趋近于0，会导致分母溢出产生NaN；另外你自定义的heavi函数不如numpy原生方法高效，还存在类型判断的潜在问题。

针对性解决方案 #

我从三个维度修改代码，解决上述问题：

1. 限制求解范围，避免m越界 #

通过设置mxstep参数限制odeint的步长，减少自变量超出给定区间的概率；同时在Ptot函数中加入判断，对负数或接近0的差值做特殊处理，避免报错。

2. 优化数值稳定性 #

替换自定义的heavi函数为numpy原生的np.heaviside，同时给sigma_sq(m)-sigma_sq(m0)设置极小阈值，避免分母为0的情况。

3. 避免除以0的情况 #

在rhs函数中判断P2(f)是否为0，防止出现除以0导致的NaN。

修改后的完整代码 #

import numpy as np
from scipy.integrate import odeint

x1 = 0.44
sigma8 = 0.82
spectral_index = 0.971
m0 = 1e13
dw = 1e-6

def sigma_sq(m):
    return sigma8**2 * m**((-2*(spectral_index+3.))/3.)

def dsdm(m):
    return (-2.*(spectral_index+3.)/3.) * m**((-2*(spectral_index+3.))/3. - 1)

def Ptot(m):
    delta_sigma = sigma_sq(m) - sigma_sq(m0)
    # 处理差值为负或接近0的情况，避免报错和数值溢出
    if isinstance(m, np.ndarray):
        mask = (delta_sigma <= 1e-12)
        result = np.zeros_like(m)
        valid_mask = ~mask
        delta_valid = delta_sigma[valid_mask]
        result[valid_mask] = (1./np.sqrt(2.*np.pi)) * (m0/m[valid_mask]) * (dw / (delta_valid)**(3./2.)) * np.abs(dsdm(m[valid_mask]))
        return result
    else:
        if delta_sigma <= 1e-12:
            return 0.0
        return (1./np.sqrt(2.*np.pi)) * (m0/m) * (dw / (delta_sigma)**(3./2.)) * np.abs(dsdm(m))

def P1(m):
    # 直接使用numpy原生阶跃函数，替代自定义实现
    return Ptot(m) * np.heaviside(m - x1*m0, 0.5)

def P2(m):
    ptot_val = Ptot(m)
    p1_val = P1(m)
    return ptot_val - p1_val

def rhs(f, m):
    p2_val = P2(f)
    # 避免除以0导致的NaN
    if isinstance(p2_val, np.ndarray):
        mask = (p2_val == 0)
        result = np.zeros_like(p2_val)
        valid_mask = ~mask
        result[valid_mask] = -P1(m[valid_mask])/p2_val[valid_mask]
        return result
    else:
        if p2_val == 0:
            return 0.0
        return -P1(m)/p2_val

# 定义求解区间
mlist2 = np.linspace(x1*m0, m0, 10000)
y0 = x1*m0
# 设置mxstep限制步长，降低越界概率
f2 = odeint(rhs, y0, mlist2, mxstep=100000)

额外说明 #

• 如果还是出现m越界的情况，可以改用更现代的求解器scipy.integrate.solve_ivp，它支持通过events参数检测m达到m0时自动停止求解。
• 代码中的1e-12是自定义的极小阈值，你可以根据计算精度需求调整。

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