Чтение hdf5-файла быстро с помощью cython и h5py

Question

Я пытаюсь ускорить функцию python3, которая принимает некоторые данные, которые представляют собой массив индексов и сохраняет их, если они соответствуют определенному критерию. Я попытался ускорить его, используя «cython -a script.py», но шейка бутылки, по-видимому, представляет собой набор данных hcpy ввода-вывода для ввода-вывода.

Я относительно не знаком с 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, а 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 массивами. В любом случае математика на массивах в памяти будет намного быстрее, чем чтение файлов.