numpy - how to replace nans by numbers in a multidimensional array in Numba?

Question

In plain python, replacing nans by numbers in-place in a numpy array is trivial. However, the following fails in when doing the same in Numba

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

How can I replace nans by zeros in a numpy array, in-place, in a Numba-compilable function? For one-dimensional one can do x[np.isnan(x)]=0 but for higher dimensions this fails as well.

Answer 1

For one-dimensional one can do x[np.isnan(x)]=0 but for higher dimensions this fails as well.

There is the clue :)

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)

Output:

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