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

Question

我正在尝试加快python3函数的速度，该函数获取一些数据，这是一个索引数组，如果它们满足某个条件，就保存它们。我试着用“cython-a”来加快速度脚本.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代码与HDF5c代码接口。所以您自己的cython代码可以直接访问它。但我想这需要更多的研究。在

您的代码正在使用Python接口来h5py，而{}不会触及这个接口。在

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数组。在任何情况下，内存中数组的运算速度都会比文件读取快得多。在