使用cython和h5py快速读取hdf5文件

Question

我正在尝试加速一个 python3 函数，该函数获取一些数据，这是一个索引数组，如果它们满足特定条件则保存它们。我试图通过使用 "cython -a script.py" 来加速它，但瓶颈似乎是 h5py I/O 切片数据集。

我对 cython 比较陌生，所以我想知道是否有任何方法可以加快速度，或者我只是受限于这里的 h5py I/O？

这是我要改进的功能：

import numpy as np
import h5py

cimport numpy as np
cimport cython
from libc.math cimport sqrt

DTYPE64 = np.int64
ctypedef np.int64_t DTYPE64_t
DTYPE32 = np.int32
ctypedef np.int32_t DTYPE32_t

@cython.boundscheck(False)
@cython.wraparound(False)
def tag_subhalo_branch(np.ndarray[DTYPE64_t] halos_z0_treeindxs,
                       np.ndarray[DTYPE64_t] tree_pindx,
                       np.ndarray[DTYPE32_t] tree_psnapnum,
                       np.ndarray[DTYPE64_t] tree_psnapid,
                       np.ndarray[DTYPE64_t] tree_hsnapid, hf,
                       int size):

    cdef int i
    cdef double radial, progen_x, progen_y, progen_z
    cdef double host_x, host_y, host_z, host_rvir
    cdef DTYPE64_t progen_indx, progen_haloid, host_id
    cdef DTYPE32_t progen_snap
    cdef int j = 0
    cdef int size_array = size
    cdef np.ndarray[DTYPE64_t] backsplash_ids = np.zeros(size_array,
                                                         dtype=DTYPE64)


    for i in range(0, size_array):
        progen_indx = tree_pindx[halos_z0_treeindxs[i]]
        if progen_indx != -1:
            progen_snap = tree_psnapnum[progen_indx]
            progen_haloid = tree_psnapid[progen_indx]

            while progen_indx != -1 and progen_snap != -1:
                # ** This is slow **
                grp = hf['Snapshots/snap_' + str('%03d' % progen_snap) + '/']
                host_id = grp['HaloCatalog'][(progen_haloid - 1), 2]
                # **

                if host_id != -1:
                    # ** This is slow **
                    progen_x = grp['HaloCatalog'][(progen_haloid - 1), 6]
                    host_x = grp['HaloCatalog'][(host_id - 1), 6]
                    progen_y = grp['HaloCatalog'][(progen_haloid - 1), 7]
                    host_y = grp['HaloCatalog'][(host_id - 1), 7]
                    progen_z = grp['HaloCatalog'][(progen_haloid - 1), 8]
                    host_z = grp['HaloCatalog'][(host_id - 1), 8]
                    # **
                    radial = 0
                    radial += (progen_x - host_x)**2
                    radial += (progen_y - host_y)**2
                    radial += (progen_z - host_z)**2
                    radial = sqrt(radial)

                    host_rvir = grp['HaloCatalog'][(host_id - 1), 24]
                    if radial <= host_rvir:
                        backsplash_ids[j] = tree_hsnapid[
                            halos_z0_treeindxs[i]]
                        j += 1
                        break

                # Find next progenitor information
                progen_indx = tree_pindx[progen_indx]
                progen_snap = tree_psnapnum[progen_indx]
                progen_haloid = tree_psnapid[progen_indx]
    return backsplash_ids

Answer 1

如此处所述：http://api.h5py.org/、h5py 使用 cython 代码与 HDF5 c 代码交互。因此，您自己的 cython 代码可能能够直接访问它。但我怀疑这需要更多的研究。

您的代码正在使用 h5py 的 Python 接口，cythonizing 不会触及它。

cython 代码最适合用于低级操作，尤其是不能表示为数组操作的迭代操作。首先研究和试验 numpy 示例。您正潜入游泳池深处的 cython。

您是否尝试过仅使用 Python 和 numpy 来改进该代码？乍一看，我看到了很多多余的 h5py 调用。

====================

您的 radial 计算访问 h5py 索引 6 次，而它可以通过 2 次访问。也许您这样写是希望 cython 可以更快地执行以下计算比 numpy?

data = grp['HaloCatalog']
progen = data[progen_haloid-1, 6:9]
host = data[host_id-1, 6:9]
radial = np.sqrt((progren-host)**2).sum(axis=1))

为什么不加载所有 data[progen_haloid-1,:] 和 data[host_id-1,:]？甚至全部data？我必须检查 h5py 何时从直接使用文件中的数组切换到 numpy 数组。无论如何，内存中数组的数学运算将比文件读取快得多。

使用cython和h5py快速读取hdf5文件

Reading hdf5 file quickly with cython and h5py

python

cython

h5py