如何在FITS数组中高效生成多个二维高斯源？

核心问题分析 #

原代码通过Python双重循环逐像素计算高斯分布，叠加n个源时需重复n次遍历整个图像，Python循环的低效率导致速度极慢。优化的核心是利用numpy向量化运算替代Python循环，将计算转移到numpy的C底层执行，大幅提升效率。

优化方案一：单源向量化叠加（简洁易读） #

利用numpy的坐标网格和广播机制，每个高斯源的计算直接针对整个图像数组，避免逐像素循环：

from astropy.io import fits
import numpy as np

# 图像尺寸
img_size = (1024, 1024)
xx, yy = img_size

# 生成一维坐标网格（ogrid比meshgrid更节省内存，自动广播为二维）
x = np.ogrid[:xx]
y = np.ogrid[:yy]

# 定义所有高斯源的参数：(中心x, 中心y, 振幅A, 标准差sigma)
gaussian_params = [
    (xx/2, yy/2, 10.0, 50.0),
    (xx/4, yy/4, 8.0, 30.0),
    (3*xx/4, 3*yy/4, 12.0, 40.0),
    # 可继续添加更多源
]

# 初始化图像数组，指定float32减少内存占用（精度允许时推荐）
final_img = np.zeros(img_size, dtype=np.float32)

# 逐个叠加高斯源，每个源的计算都是向量化操作
for xc, yc, A, sigma in gaussian_params:
    # 利用numpy广播自动扩展维度，一次性计算整个图像的高斯值
    gaussian = A * np.exp(-((x - xc)/sigma)**2 - ((y - yc)/sigma)**2)
    final_img += gaussian

# 保存为FITS文件
fits.writeto('multi_gaussian_model.fits', final_img, overwrite=True)

优化方案二：全向量化批量计算（极致性能） #

当高斯源数量较多时，可将所有参数转为numpy数组，利用三维广播一次性计算所有源的贡献，再求和叠加，完全消除Python循环：

from astropy.io import fits
import numpy as np

img_size = (1024, 1024)
xx, yy = img_size

# 生成坐标网格
x = np.ogrid[:xx]
y = np.ogrid[:yy]

# 将所有源的参数转为numpy数组
xc_list = np.array([xx/2, xx/4, 3*xx/4])
yc_list = np.array([yy/2, yy/4, 3*yy/4])
A_list = np.array([10.0, 8.0, 12.0])
sigma_list = np.array([50.0, 30.0, 40.0])

# 扩展维度以支持三维广播：(源数量, 图像高度, 图像宽度)
x_exp = x[np.newaxis, :, np.newaxis]
y_exp = y[np.newaxis, np.newaxis, :]
xc_exp = xc_list[:, np.newaxis, np.newaxis]
yc_exp = yc_list[:, np.newaxis, np.newaxis]
A_exp = A_list[:, np.newaxis, np.newaxis]
sigma_exp = sigma_list[:, np.newaxis, np.newaxis]

# 一次性计算所有高斯源，沿源维度求和得到最终图像
final_img = np.sum(
    A_exp * np.exp(-((x_exp - xc_exp)/sigma_exp)**2 - ((y_exp - yc_exp)/sigma_exp)**2),
    axis=0
)

fits.writeto('multi_gaussian_model.fits', final_img, overwrite=True)

关键优化点说明 #

1. 向量化运算：numpy的底层运算由C实现，比Python循环快几十到上百倍，尤其适合大尺寸图像和多源场景。
2. 内存优化：使用np.ogrid生成一维坐标数组，通过广播自动扩展为二维，比meshgrid节省一半内存。
3. 数据类型：指定dtype=np.float32可将内存占用减少一半，同时不损失天文图像常用的精度。

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