使用For循环绘制演化恒星合成光谱（SSP）的错误排查

我有一批不同恒星类型的FITS文件，每个文件包含流量（flux）和波长（wavelength）列。已计算出以下变量：

• mainseq_gyr：主序星寿命列表（单位：Gyr）
• weighted_spect：各光谱型相对总样本的质量权重比列表
• flux_per_col：每个光谱型的流量数组列表（每个数组含4771个元素）
• wavelength：所有文件共用的统一波长范围数组

尝试绘制不同时间步长下恒星族群的SED，但当前代码只会重复绘制同一张图，最终仅显示最后一个时间步长的结果，请求修正代码实现正确绘图。

当前代码如下：

timesteps  = [0.01, 0.1, 0.2, 0.5, 1, 2, 5, 10] #in Gyr
fig, ax = plt.subplots()
labels = ["0.01 Gyr", "0.1 Gyr","0.2 Gyr", "0.5 Gyr","1 Gyr","2 Gyr", "5 Gyr", "10 Gyr"]
for i,t in enumerate(timesteps):
    SSP = []
    for tau,w,fl in zip(mainseq_gyr,weighted_spect,flux_per_col):     
        if tau > t:
            SSP.append(fl*w)
            #print(fl*w)
    fssp = sum(SSP)
    ax.plot(wavelength,fssp,label = labels[i])
    
ax.legend()
ax.set_xlabel('$\lambda(A^{\circ})$')
ax.set_ylabel('$Flux  f_{v}$')

问题核心在于空列表求和的逻辑漏洞：当时间步长t大于所有恒星主序寿命tau时，SSP为空列表，此时sum(SSP)的返回值不符合数组运算预期，会导致后续绘图覆盖之前的结果。同时，原循环逐个处理数组的方式效率较低，也容易引发维度匹配问题。

修正后的代码如下：

import numpy as np
import matplotlib.pyplot as plt

timesteps  = [0.01, 0.1, 0.2, 0.5, 1, 2, 5, 10] #in Gyr
fig, ax = plt.subplots()
labels = ["0.01 Gyr", "0.1 Gyr","0.2 Gyr", "0.5 Gyr","1 Gyr","2 Gyr", "5 Gyr", "10 Gyr"]

# 将列表转为numpy数组，方便批量运算
flux_per_col_np = np.array(flux_per_col)
weighted_spect_np = np.array(weighted_spect)
mainseq_gyr_np = np.array(mainseq_gyr)

for i,t in enumerate(timesteps):
    # 筛选寿命大于当前时间步长的恒星
    mask = mainseq_gyr_np > t
    if np.any(mask):
        # 权重数组添加维度，和流量数组广播相乘后求和
        fssp = np.sum(flux_per_col_np[mask] * weighted_spect_np[mask, np.newaxis], axis=0)
    else:
        # 无符合条件恒星时，生成全零数组
        fssp = np.zeros_like(wavelength)
    ax.plot(wavelength, fssp, label=labels[i])

ax.legend()
ax.set_xlabel('$\lambda(A^{\circ})$')
ax.set_ylabel('$Flux  f_{v}$')
plt.show()

关键修正点： #

• 改用numpy批量运算：替代循环逐个处理数组，提升效率的同时避免列表求和的潜在问题
• 显式处理空筛选场景：无符合条件恒星时生成与波长等长的全零数组，保证绘图逻辑统一
• 维度匹配处理：通过weighted_spect_np[mask, np.newaxis]让权重数组和流量数组正确广播相乘，确保维度一致

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