CODEWITHSUNDEEP
arraysnumpyboolean-logicnumbaboolean-operations
stagermane
stagermane

Reputation: 1063

Numba nopython mode boolean indexing

I'm working with arrays in numba that sometimes return inf or nan as an entry, which I then need to adjust to zero. I'm working in nopython mode, so I want to know if there's a way to use boolean indexing (or multiplication) to adjust any nan or inf elements in the array to zero (array shape is (vertex_number, 1)).

Obviously I can just loop over the array like this:

from numpy import isnan, isinf, array, zeros, nan, inf, ones

from numba import jit

v = ones((10, 1))

v[0] = nan
v[-1] = inf
print(v)

@jit(nopython=True, nogil=True)
def reciprocal_v(v):
    u = 1 / v
    for i in range(len(u)): # get rid of zero length reciprocals...
        if (isinf(u[i])[0] == True) | (isnan(u[i])[0] == True):
            u[i] = 0.0
    return u
u = reciprocal_v(v)
print(u)

I tried boolean multiplication, but inf*0 returns nan, and nan*0 returns nan

@jit(nopython=True)
def numba_isnan(x): return isnan(x)

@jit(nopython=True)
def numba_isinf(x): return isinf(x)

@jit(nopython=True, nogil=True)
def reciprocal_v_bool(v):
    u = 1 / v
    u = u * ~(numba_isnan(u) | numba_isinf(u))
    return u

Upvotes: 1

Views: 1408

Answers (1)

JoshAdel
JoshAdel

Reputation: 68682

This isn't an issue with numba, but instead it's the result of how numpy handles the multiplication of nan and inf (and numba follows the convention). If you remove the jit decorator, you can see that the un-jitted function returns the same result. This behavior is defined by a standard like the IEEE 754.

I would just stick with the solution you've already come up with that works.

Also as a general note, you don't need functions like:

@jit(nopython=True)
def numba_isnan(x): return isnan(x)

Just import numpy as np and then use np.isnan() in a numba function. Numba will recognize the numpy function and replace it with the correct form.

Upvotes: 2

Related Questions