运行热方程Python代码触发IndexError，求解决方案

问题原因 #

报错核心是数组索引必须为整数类型，代码中T[t/dt,:]里的t/dt计算结果是浮点数（比如0.01/0.0005=20.0），即使数值是整数，numpy也不允许用浮点数索引数组，因此抛出IndexError。

解决方法 #

将浮点数时间值转换为对应的整数时间步索引，同时处理浮点数精度问题，避免索引超出数组范围。

修改后的代码示例 #

# solve the heat equation for a metal rod as one end is kept at 100 °C 
# and the other at 0 °C 
# u(0,t) = 100 (x=0)
# u(L,t) = 0 (x=L)

import numpy as np
import matplotlib.pyplot as plt

dt = 0.0005
dy = 0.0005
k = 10**(-4)
y_max = 0.04
t_max = 1
T0 = 100

def FTCS(dt,dy,t_max,y_max,k,T0):
    s = k*dt/dy**2
    y = np.arange(0,y_max+dy,dy) 
    t = np.arange(0,t_max+dt,dt)
    r = len(t)
    c = len(y)
    T = np.zeros([r,c])
    T[:,0] = T0
    for n in range(0,r-1):
        for j in range(1,c-1):
            T[n+1,j] = T[n,j] + s*(T[n,j-1] - 2*T[n,j] + T[n,j+1]) 
    return y,T,r,s

y,T,r,s = FTCS(dt,dy,t_max,y_max,k,T0)

plot_times = np.arange(0.01,1.0,0.01)
for t in plot_times:
    # 将时间转换为整数索引，用round处理浮点数精度误差
    time_idx = int(np.round(t / dt))
    # 确保索引不超出数组行数范围
    if time_idx < r:
        plt.plot(y, T[time_idx, :])
plt.show()

额外优化方案 #

如果想彻底避免浮点数计算的精度问题，可以直接基于时间步索引循环，这样更可靠：

# 计算起始和结束的时间步索引
start_idx = int(np.round(0.01 / dt))
end_idx = int(np.round(1.0 / dt))
# 按0.01时间间隔对应的步长循环
step = int(np.round(0.01 / dt))

for time_idx in range(start_idx, end_idx, step):
    plt.plot(y, T[time_idx, :])
plt.show()

关键说明 #

• T数组的每一行对应一个时间步，行索引n对应的实际时间是n*dt
• 浮点数运算可能存在精度偏差（比如0.03可能被存储为0.029999999999999998），用np.round()可以确保得到正确的整数索引
• 增加if time_idx < r的判断，防止因计算误差导致索引超出数组范围

内容的提问来源于stack exchange，提问作者Freya the Goddess