基于Python使用多份FITS文件创建HDF5数据立方体

哇，处理1500个3800×3800的FITS文件确实是个实打实的挑战——单帧就有近1500万像素，1500帧堆起来的数据量相当惊人，难怪单循环写入会慢到每帧30分钟。我在处理天文大数据时踩过不少类似的坑，分享几个亲测有效的优化思路，应该能帮你把速度提上去：

1. 给HDF5设置合理的分块参数 #

你已经创建了空的HDF5数组，但如果没指定chunks参数，默认是整数组连续存储，写入时磁盘IO开销极大。建议创建数据集时根据单帧大小设置分块，比如每5-10个帧作为一个块：

import h5py
import numpy as np

# 创建带分块的HDF5数据集
with h5py.File('cube.h5', 'w') as h5f:
    h5f.create_dataset(
        'data_cube',
        shape=(3800, 3800, 1500),
        dtype=np.float32,  # 根据你的FITS数据类型调整
        chunks=(3800, 3800, 10),  # 每10帧一个块
        shuffle=True,  # 打乱数据提升压缩效率
        compression='lzf'  # 轻量压缩，几乎不损失速度还能减磁盘占用
    )

分块后写入时，磁盘会按块读写，而不是每次操作整个超大数组，能显著降低IO耗时。

2. 批量读取+批量写入代替单帧循环 #

单帧读取写入的开销主要来自重复打开/关闭FITS文件、内存拷贝和磁盘IO。可以一次读取多个FITS文件（比如10个），然后一次性写入HDF5的对应切片：

from astropy.io import fits
import glob

fits_files = sorted(glob.glob('*.fits'))
batch_size = 10

with h5py.File('cube.h5', 'r+') as h5f:
    cube_dset = h5f['data_cube']
    for i in range(0, len(fits_files), batch_size):
        batch_files = fits_files[i:i+batch_size]
        # 批量读取当前批次的FITS数据
        batch_data = []
        for fpath in batch_files:
            with fits.open(fpath, memmap=True) as hdul:
                # 用memmap避免把整帧加载到内存，节省资源
                batch_data.append(hdul[0].data)
        # 把批次数据堆叠成(3800,3800,batch_size)的数组
        batch_arr = np.stack(batch_data, axis=2)
        # 写入HDF5的对应位置
        cube_dset[:, :, i:i+batch_size] = batch_arr
        print(f'完成批次 {i//batch_size + 1}/{len(fits_files)//batch_size + 1}')

这样能减少文件IO的次数，提升整体效率。

3. 用多进程并行读取FITS文件 #

单线程读取1500个文件太慢，可以用多进程并行读取数据，再由主线程统一写入HDF5（注意HDF5不支持多进程同时写入，所以写入必须单线程）：

from concurrent.futures import ProcessPoolExecutor

def read_fits(fpath):
    with fits.open(fpath, memmap=True) as hdul:
        return hdul[0].data

with h5py.File('cube.h5', 'r+') as h5f:
    cube_dset = h5f['data_cube']
    # 用进程池并行读取，进程数根据你的CPU核心数调整
    with ProcessPoolExecutor(max_workers=4) as executor:
        # 按批次提交任务，避免一次性加载太多数据到内存
        for i in range(0, len(fits_files), batch_size):
            batch_futures = [executor.submit(read_fits, f) for f in fits_files[i:i+batch_size]]
            batch_data = [future.result() for future in batch_futures]
            batch_arr = np.stack(batch_data, axis=2)
            cube_dset[:, :, i:i+batch_size] = batch_arr
            print(f'完成批次 {i//batch_size + 1}')

并行读取能充分利用CPU资源，把读取的时间压缩下来，这对机械硬盘的提升尤为明显。

4. 检查磁盘IO瓶颈 #

如果上述优化后速度还是慢，大概率是磁盘拖了后腿：

• 尽量把FITS源文件和HDF5目标文件都放到SSD上，SSD的随机读写速度比机械硬盘快几十倍；
• 避免同时运行其他磁盘密集型任务（比如备份、下载），减少IO竞争；
• 如果是网络存储（NAS），尽量用高速网络连接（比如10Gbps以太网），或者把文件先拷贝到本地磁盘再处理。

先从调整HDF5的分块和压缩参数开始，这是最容易实施的优化；然后尝试批量读取和并行处理，应该能把单次迭代的时间从30分钟降到几分钟甚至更短。记得每次优化后先测试一小部分数据（比如10个文件），确认效果再全量运行哦！

内容的提问来源于stack exchange，提问作者Alejandro Serrano Borlaff