使用多个 FITS 文件在 Python 中创建 HDF5 数据立方体

Question

我目前正在努力解决一个问题。我有 1500 个包含 3800 x 3800 数组的拟合文件。我的 objective 是用它们创建一个 HDF5 数据立方体。不幸的是，我无法提供合适的文件（由于存储问题）。到目前为止，我已经能够通过执行以下操作创建具有所需形状（3800、3800、1500）的空 HDF5 数组：

import h5py    
from astropy.io import fits
import glob

outname = "test.hdf5"
NAXIS1 = 3800
NAXIS2 = 3800
NAXIS3 = 1500
f = h5py.File(outname,'w')
dataset = f.create_dataset("DataCube",
                           (NAXIS1,NAXIS2,NAXIS3),
                           dtype=np.float32)
f.close()

但是我在尝试从 fits 文件中写入数组时遇到了问题，因为以下 for 循环中的每个元素至少需要 30 分钟：

f = h5py.File(outname,'r+')
# This is the actual line, but I will replace it by random noise
# in order to make the example reproducible.
# fitslist = glob.glob("*fits") # They are 1500 fits files

for i in range(NAXIS3): 
    # Again, I replace the real data with noise.
    # hdul = fits.open(fitslist[i])
    # file['DataCube'][:,:,i] = hdul[0].data     
    data = np.random.normal(0,1,(dim0,dim1))    
    file['DataCube'][:,:,i] = data
f.close()

有没有更好的方法来构建一个由 N 个切片组成的 3D 数据立方体，这些切片已经存储在 N 个适合的文件中？我期待在光盘中创建 HDF5 文件后，写入它会非常快，但事实并非如此。

非常感谢您的帮助。

编辑 1：我测试了 astrofrog 提出的修改，效果非常好。现在的表现已经很不错了。除此之外，我将几个适合的文件（~50）存储到一个临时的 numpy 数组中，以减少我写入 hdf5 文件的次数。现在代码如下所示：

NAXIS1 = len(fitslist)
NAXIS2 = fits_0[ext].header['NAXIS1']
NAXIS3 = fits_0[ext].header['NAXIS2']
shape_array = (NAXIS2, NAXIS3)
print(shape_array)

f = h5py_cache.File(outname, 'w', chunk_cache_mem_size=3000*1024**2,
                    libver='latest')

dataset = f.create_dataset("/x", (NAXIS1, NAXIS2, NAXIS3),
                           dtype=np.float32)

cache_size = 50
cache_array = np.empty(shape=(cache_size, NAXIS2, NAXIS3))
j = 0
for i in tqdm(range(len(fitslist))):
    print(fitslist[i])
    hdul = fits.getdata(fitslist[i], ext)
    cache_array[j:j+1, :, :] = hdul

    if ((i % cache_size == 0) & (i != 0)):
        print("Writing to disc")
        f['/x'][i-cache_size+1:i+1, :, :] = cache_array
        j = 0
    if (i % 100 == 0):
        print("collecting garbage")
        gc.collect()
    j = j + 1
f.close()

我的问题是：还有更多的 pythonic 方法可以做到这一点吗？我不确定这是使用 h5py 写入文件的最有效方式，还是有更好的方式从 fits 读取到 numpy，然后再读取到 hdf5。

Answer 1

我认为问题可能是维度的顺序应该是 NAXIS3、NAXIS2、NAXIS1（目前我认为它在数组上的跨度非常低效）。我也只会在最后将数组添加到 HDF5 文件中：

import glob

import h5py
import numpy as np
from astropy.io import fits

fitslist = glob.glob("*.fits")

NAXIS1 = 3800
NAXIS2 = 3800
NAXIS3 = 1000

array = np.zeros((NAXIS3, NAXIS2, NAXIS1), dtype=np.float32)

for i in range(NAXIS3):
    hdul = fits.open(fitslist[i], memmap=False)
    array[i, :, :] = hdul[0].data

f = h5py.File('test.hdf5', 'w')
f.create_dataset("DataCube", data=array)
f.close()

如果需要NAXIS1,NAXIS2,NAXIS3顺序的数组，在最后转置即可。

使用多个 FITS 文件在 Python 中创建 HDF5 数据立方体

Creating a HDF5 datacube in Python with multiple FITS files

python

hdf5

data-cube

fits

astropy