Оптимизация булевой индексации Cython

Question

Каков наилучший способ преобразовать следующий код в cython

Учитывая следующий пример:

#setup example data Z and A
Z = np.random.randn(10,10)
A = np.random.randn(10,10)
A[0,1] = np.nan
A[1,3] = np.nan
A[5,3] = np.nan
A[3,5] = np.nan

B = np.isnan(A).transpose()

C = Z[B * B.transpose()]

Я хочу оптимизировать определение типа np.ndarray B в приведенном выше примере и оптимизировать создание ndarray C,

Я попытался использовать настройку B для uint8 и python и C++ bools.

cdef np.ndarray[np.uint8_t, ndim=2, cast=True] однако это приводит к небольшому ускорению или его отсутствию

а также

cdef np.ndarray[bool, ndim=2, cast=True]

где bool либо from cpython cimport bool или же from libcpp cimport bool в обоих случаях приведенный выше код выдаст ошибку.

Answer 1

Правильный способ создать буфер, который будет принимать np.nan значения использует np.float_t или же np.double_t, Если вы попытаетесь использовать целочисленный буфер, появится следующая ошибка:

ValueError: cannot convert float Nan to integer

Затем вы можете использовать что-то вроде:

cdef np.ndarray[np.double_t, ndim=2] A, Z

Z = np.random.randn(10,10)
A = np.random.randn(10,10)
A[0,1] = np.nan
A[1,3] = np.nan
A[5,3] = np.nan
A[3,5] = np.nan
B = np.isnan(A).transpose()
C = Z[B * B.transpose()]