Python遍历FITS数据立方实现逐像素光谱高斯拟合

核心调整思路 #

• 把所有和像素位置无关的公共计算（静止波长校准、通道索引筛选、拟合x轴数组）提前移到循环外，避免重复计算浪费算力
• 将单像素拟合逻辑封装为独立函数，加入异常捕获机制，避免坏像素、低信噪比像素拟合失败导致整个任务中断
• 提前初始化和空间维度尺寸匹配的结果存储数组，遍历所有空间像素坐标完成拟合后直接写入对应位置
• 优先使用已经验证过的中心像素拟合结果作为全局初始参数，比逐像素计算初始值的拟合稳定性高很多

完整可运行代码 #

import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from spectral_cube import SpectralCube
from scipy.optimize import curve_fit

# -------------------------- 1. 全局公共数据加载与预计算（所有像素共用，仅执行1次） --------------------------
c = 2.998 * 10**5  # 光速，单位km/s
ngc2617 = fits.open('NGC2617.fits')[1]
cube2617 = SpectralCube.read(ngc2617)
data2617 = ngc2617.data
# 统一替换全数据立方的NaN值为0，避免后续计算报错
data2617[np.isnan(data2617)] = 0
nz, ny, nx = np.shape(data2617)

# 波长校准
obs_2617 = cube2617.spectral_axis
z_2617 = 0.014326  # 红移值，来自simbad
rest_2617 = np.asarray(obs_2617 / (1+z_2617))

# 预计算Hα线区、连续谱的通道索引，所有像素通用
halpha_chans = np.where((rest_2617 > 6365) & (rest_2617 < 6715))[0]
cont_chans = np.where((rest_2617 > 5500) & (rest_2617 < 5510))[0]
x_halpha = rest_2617[halpha_chans]  # 拟合用的波长x轴，全局通用

# 高斯函数与组合拟合模型定义（全局通用）
def Gauss(x, A, x0, fwhm):
    sigma = (fwhm * x0) / (2.355 * c)
    return A * np.exp(-(x - x0) ** 2 / (2 * sigma ** 2))

def func(x, A0, A1, x0, fwhm, fwhm1, cont):
    return cont + Gauss(x, A0, x0 - 14.75, fwhm) + Gauss(x, A1, x0, fwhm) + Gauss(x, A0*3, x0 + 20.66, fwhm) + Gauss(x, A1, x0, fwhm1)

# 先用已验证的中心像素(172,142)拟合得到可靠初始值，作为全局拟合初始p0
flux_center = data2617[:, 172, 142]
halphaflux_center = flux_center[halpha_chans]
cont_center = np.average(flux_center[cont_chans])
A0_init = np.max(halphaflux_center[0:150]) - cont_center
A1_init = np.max(halphaflux_center) - cont_center
x0_init = x_halpha[np.argmax(halphaflux_center)]
fwhm_init = (x_halpha[np.argmax(halphaflux_center)] - x_halpha[np.argmin(np.abs(np.max(halphaflux_center)/2 - halphaflux_center))])*2
fwhm1_init = (x_halpha[np.argmax(halphaflux_center)] - x_halpha[np.argmin(np.abs(np.max(halphaflux_center)/2 - halphaflux_center[0:150]))])*2
p0_global = [A0_init, A1_init, x0_init, fwhm_init, fwhm1_init, cont_center]

# -------------------------- 2. 单像素拟合函数封装 --------------------------
def fit_single_pixel(flux):
    """
    输入单像素沿光谱轴的通量数组，返回拟合参数数组，拟合失败返回全NaN
    返回参数顺序: A0, A1, x0, fwhm, fwhm1, cont
    """
    try:
        y_halpha = flux[halpha_chans]
        cont_val = np.average(flux[cont_chans])
        # 针对当前像素微调初始值里的连续谱和峰值，进一步提升稳定性
        p0 = p0_global.copy()
        p0[-1] = cont_val
        p0[1] = np.max(y_halpha) - cont_val
        popt, pcov = curve_fit(func, x_halpha, y_halpha, p0=p0, maxfev=5000)
        return popt
    except (RuntimeError, ValueError):
        # 拟合不收敛、数据异常时返回空值
        return np.full(6, np.nan)

# -------------------------- 3. 遍历全像素批量拟合 --------------------------
# 提前初始化结果数组，维度为(参数个数, ny, nx)，如果只需要特定物理量也可以单独建数组
# 比如只存Hα流量就建尺寸为(ny,nx)的数组即可
fit_results = np.full((6, ny, nx), np.nan)

# 如果只需要拟合10*10的子区域，把下面循环的y、x范围改成对应切片即可，比如range(167,177)就是中心附近10*10
for y in range(ny):
    print(f"正在拟合第{y+1}/{ny}行像素")
    for x in range(nx):
        flux_pixel = data2617[:, y, x]
        popt = fit_single_pixel(flux_pixel)
        fit_results[:, y, x] = popt

# -------------------------- 4. 结果验证（可选，抽单个像素画图确认拟合效果） --------------------------
test_y, test_x = 172, 142  # 用中心像素测试
test_flux = data2617[:, test_y, test_x][halpha_chans]
test_popt = fit_results[:, test_y, test_x]
test_fit = func(x_halpha, *test_popt)

plt.figure(figsize=[10,8])
plt.subplot(2,1,1)
plt.plot(x_halpha, test_flux, '--', label='观测数据')
plt.plot(x_halpha, test_fit, '-', label='拟合结果')
plt.xlabel('静止波长 [$\AA$]')
plt.ylabel(r'流量 [10$^{-20}$ $erg * s^{-1}$ * $cm^{-2}$ * $\AA^{-1}$]')
plt.legend()
plt.grid(True)

plt.subplot(2,1,2)
plt.scatter(x_halpha, test_flux - test_fit, marker='o', label='残差')
plt.axhline(0, color='r')
plt.xlabel('静止波长 [$\AA$]')
plt.ylabel('残差')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()

使用提示 #

• 如果只需要拟合指定小区域（比如提到的10×10范围），直接把双层循环里的range(ny)、range(nx)替换为对应坐标区间即可，比如要拟合中心附近10×10，就写for y in range(167,177)和for x in range(137,147)，运行速度会快很多
• 如果后续要处理整幅大尺寸立方，可以把循环替换为joblib并行计算，核心逻辑不用改，只需要把单层循环包装为并行任务即可
• 可以根据自身需求修改结果存储的内容，比如直接计算Hα积分流量、NII/Hα线比、速度弥散等物理量存到结果数组里，不用存全部拟合参数

内容的提问来源于stack exchange，提问作者Lacey A-P