如何在Numba中用多维数组中的数字替换NaN？

问题描述

在普通的Python中，在数字数组中就地用数字替换NaN是微不足道的。但是，在Numba中执行相同操作时，以下操作失败

@jit(nopython=True)
def dostuff():
    x = np.array([[1,np.nan,3]]);
    np.nan_to_num(x,copy=False);
dostuff()

如何在Numba可编译函数中就地用数字数组中的零替换NaN？对于一维，可以x[np.isnan(x)]=0，但对于更高的维，这也失败了。

---

解决方案

对于一维，可以做x[np.isnan(X)]=0，但对于更高的 维度此操作也失败。

有线索：)

import numpy as np
from numba import jit

@jit(nopython=True)
def dostuff(x):
    shape = x.shape
    x = x.ravel()
    x[np.isnan(x)] = 0
    x = x.reshape(shape)
    return x

a = np.array([[1,np.nan,3], [np.nan, 2, 6]])
dostuff(a)
print(a)

输出：

[[1. 0. 3.]
 [0. 2. 6.]]