RK4方法适配数组输出函数报错：序列无法与浮点数相乘

问题重现 #

编写适配数组输出的RK4方法时，调用fb2函数（返回数组形式的导数），运行时在k2行持续报错can't multiply sequence by non-int of type 'float'，尝试将步长h转为int类型也无法解决。

错误根源 #

1. 返回类型不匹配：fb2返回的是Python原生列表，列表不支持直接与浮点数进行元素级乘法（比如0.5*k1*h中，k1是列表，无法和float运算）
2. 错误修改步长类型：强行把h转为int会彻底破坏RK4的精度，且完全不是报错的核心原因
3. 变量名冲突+初始值赋值错误：代码中先把初始条件的x值赋值给x0，后又用时间起点覆盖x0；同时初始化y_val[:,0]时错误地只用了y0，导致x的初始值丢失

修复步骤 #

1. 修改fb2，让它返回numpy数组而非列表，支持浮点数元素运算
2. 删除h=int(h)这行错误代码，步长必须保留为浮点数
3. 修正变量名冲突，把初始条件的变量和时间起点变量区分开
4. 正确初始化y_val的初始值，直接用传入的初始条件数组R

修复后的完整代码 #

import numpy as np
import matplotlib.pyplot as plt

def fb2(t, J):
    x = J[0]
    y = J[1]
    dxdt = 0.25*y - x
    dydt = 3*x - y
    # 返回numpy数组而非列表，支持元素级运算
    return np.array([dxdt, dydt])

def odeRK4(function, tspan, R, h, *args):
    # 区分初始条件和时间起点，避免变量名冲突
    r0 = np.array(R)
    t_start = tspan[0]
    t_end = tspan[1]

    # 处理步长逻辑
    if h is None:
        h = 0.01 * (t_end - t_start)
    elif h > (t_end - t_start):
        h = 0.01 * (t_end - t_start)

    # 生成时间序列
    t_values = np.arange(t_start, t_end + h, h)
    N = len(t_values)
    
    # 初始化结果数组：2行N列，对应x和y的所有时间点值
    y_val = np.zeros([2, N])
    y_val[:, 0] = r0  # 正确设置初始条件

    for i in range(1, N):
        t_prev = t_values[i-1]
        y_prev = y_val[:, i-1]
        
        # RK4核心计算，现在所有变量都是numpy数组，支持元素运算
        k1 = function(t_prev, y_prev, *args)
        k2 = function(t_prev + 0.5*h, y_prev + 0.5*k1*h, *args)
        k3 = function(t_prev + 0.5*h, y_prev + 0.5*k2*h, *args)
        k4 = function(t_prev + h, y_prev + k3*h, *args)
        
        y_val[:, i] = y_prev + h*(1/6)*(k1 + 2*k2 + 2*k3 + k4)

    # 合并时间和结果并打印
    result = np.column_stack((t_values, y_val.T))  # 转置y_val让列对应x,y
    print('At time t, Y= (t on left, x on middle, y on right)')
    print(result)
    
    # 绘制x和y的时间曲线
    plt.plot(t_values, y_val[0, :], label='x(t)')
    plt.plot(t_values, y_val[1, :], label='y(t)')
    plt.legend()
    plt.xlabel('t')
    plt.ylabel('Value')
    plt.show()

print('for 3B:')
odeRK4(fb2, [0, 20], [1, 1], None)

额外说明 #

• 修复后的代码不仅解决了报错问题，还修正了初始值赋值错误的bug，确保x和y的初始条件都被正确设置
• 绘图部分添加了图例和标签，更直观展示结果
• 所有数组运算都基于numpy，保证了数值计算的效率和正确性

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