Инициализировать numba guvectorize выходной массив до нуля

Я пытаюсь создать функцию guvectorize, которая перебирает данные и по мере необходимости меняет значение ячеек в выходной матрице. Беда в том, что когда nopython=True, я не могу получить выходную матрицу для инициализации в ноль, и в результате вывод содержит бессмысленные числа.

import numpy as np
from numba import guvectorize

@guvectorize(['void(float64[:], float64[:], int8[:])'], '(n),(n)->(n,n)', nopython=True)
def test(array_1, array_2, res):
    res[:] = 0
    for i in range(array_1.size):
        if array_1[i] > 1:
            res[i] += 1

a = np.float64(np.random.randint(0, 2, (6)))
b = np.float64(np.random.randint(0, 2, (6)))
c = test(a, b)

print(c)

Выход:

[[   0    3 -126    4    0    0]
 [   0    0    0   53   26    2]
 [   0    0    0    0    0   53]
 [   0    2    0    0    0    0]
 [   0   59   33    2    0    0]
 [   0    0   56   91   58   93]]

Это дает правильный вывод, когда nopython=False, но тогда я теряю огромную скорость обработки.

Я могу исправить проблему с помощью итеративной инициализации в ноль в начале кода (вместо res[:]):

for i in range(unique_years.size):
    for j in range(window_sizes.size):
        res[i, j] = 0

но это кажется непифоническим и, возможно, медленнее, чем простой актерский состав. Использование res[:,:] вызывает ошибку приведения.

0 ответов

Другие вопросы по тегам