Reading hdf5 file quickly with cython and h5py
<p>I'm trying to speed up a python3 function that takes some data, which is an array of indexes and saves them if they meet a certain criterion. I have tried to speed it up by using "cython -a script.py", but the bottle neck seems to be the h5py I/O slicing datasets.</p>

<p>I'm relatively new to cython, so I was wondering whether there is anyway to speed this up or am I just limited by the h5py I/O here?</p>

<p>Here is the function I'm trying to improve:</p>

<pre><code>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 &lt;= 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
</code></pre>

<p>As described here: <a href="http://api.h5py.org/" rel="nofollow">http://api.h5py.org/</a>, <code>h5py</code> uses <code>cython</code> code to interface with the <code>HDF5</code> <code>c</code> code.  So your own <code>cython</code> code might be able to access that directly.  But I suspect that will require a lot more study.</p>

<p>Your code is using the Python interface to <code>h5py</code>, and <code>cythonizing</code> isn't going to touch that.</p>

<p><code>cython</code> code is best used for low level actions, especially iterative things that can't be expressed as array operations.  Study and experiment with the <code>numpy</code> examples first.  You are diving into <code>cython</code> at the deep end of the pool.</p>

<p>Have you tried to improve that code just with Python and numpy?  Just at glance I'm seeing a lot of redundant <code>h5py</code> calls.</p>

<p>====================</p>

<p>Your <code>radial</code> calculation accesses the <code>h5py</code> indexing 6 times when it could get by with 2.  Maybe you wrote it that way in hopes that <code>cython</code> would preform the following calculation faster than numpy?</p>

<pre><code>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))
</code></pre>

<p>Why not load all <code>data[progen_haloid-1,:]</code> and <code>data[host_id-1,:]</code>?  Even all of <code>data</code>?  I'd have to review when <code>h5py</code> switches from working directly with the arrays on the file and when they become <code>numpy</code> arrays.  In any case, math on arrays in memory will be a lot faster than file reads.</p>

python cython h5py