Как сравнить скалярное значение с массивом в Numba (Python)? - PullRequest
Новая
       8

Как сравнить скалярное значение с массивом в Numba (Python)?

0 голосов
/ 28 февраля 2020

Как сравнить скалярное значение с массивом в Numba (Python)? У меня есть функция, которая должна принимать 1-D и 2-D массивы в качестве параметров. В рамках функции я должен проверить, является ли фактическое значение для каждой строки не NaN. Я делаю это с numpy функцией np.isnan. Эта функция предоставляет скаляры для 1-D массивов и массивы для 2-D массивов. К сожалению, Numba не может решить эту проблему. Как я могу проверить, имеет ли строка массива значения NaN в 1-D или 2-D массивах? Нужно ли писать 3 функции: главную функцию, а затем делегировать функцию для 1-D или 2-D массивов. 

import numpy as np
import numba as nb
@nb.jit([nb.float64[:](nb.float64[:]), nb.float64[:,:](nb.float64[:,:])], nopython=True)
def is_not_nan_in_array_or_scalar(x):
    result = np.full_like(x, np.NaN)
    is_not_nan = ~np.isnan(x) if x.ndim==1 else ~np.isnan(x.flatten())
    for i in range(0, len(x)):
        if is_not_nan[i]:
            # do something...
            result[i] = is_not_nan[i]
    return result

print(is_not_nan_in_array_or_scalar(np.array([1,2.5,np.NaN])))
print(is_not_nan_in_array_or_scalar(np.array([[1],[2.5],[np.NaN]])))

Это должно быть результатом: 

[ 1.  1. nan]
[[ 1.]
 [ 1.]
 [nan]]

Это ошибка с использованием Numba: 

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Cannot unify array(bool, 2d, C) and array(bool, 1d, C) for '$phi50.1', defined at <ipython-input-6-2bd50d9143ab> (6)

File "<ipython-input-6-2bd50d9143ab>", line 6:
def is_not_nan_in_array_or_scalar(x):
    <source elided>
    result = np.full_like(x, np.NaN)
    is_not_nan = ~np.isnan(x) if x.ndim==1 else ~np.isnan(x.flatten())
    ^

[1] During: typing of assignment at <ipython-input-6-2bd50d9143ab> (6)

File "<ipython-input-6-2bd50d9143ab>", line 6:
def is_not_nan_in_array_or_scalar(x):
    <source elided>
    result = np.full_like(x, np.NaN)
    is_not_nan = ~np.isnan(x) if x.ndim==1 else ~np.isnan(x.flatten())
    ^
...