Unpacking np.where index alternative

Numba Support: How do I do ...?
1

Hello all,

I was wondering if there is any easier way to index a matrix using the tuples coming from np.where, I have came up with a quick fix, but maybe there is a better alternative.

import numpy as np

from numba import jit

@jit(nopython=True)

def find_higher(value, array_1):

    ix_higher = np.where(array_1 < value)

    return array_1[ix_higher]

@jit(nopython=True)

def find_higher_fixed(value, array_1):

    ix_higher = np.where(array_1 < value)

    ixs = len(ix_higher[0])

    answer = np.zeros(ixs)

    if ixs > 0:

        answer = np.zeros(ixs)

        for ix in range(ixs):

            answer[ix] = array_1[ix_higher[0][ix], ix_higher[1][ix]]

    return answer

a = np.ones(shape=(3, 3))

a[0, 0] = 2

a[1, 2] = 2

print(a)

b = find_higher(2, a)  # Error getitem(array(float64, 2d, C), UniTuple(array(int64, 1d, C) x 2))

print(b)

b = find_higher_fixed(

    2, a

)  # Works, but is there a better option?

print(b)

Any suggestions are welcomed :slight_smile:

2

hi @DTKx , do you really need to use np.where? Numba works very well with loops, and I personally find it easier to write and read loops than np.where. Plus, the loop will probably be faster (because it avoids allocating intermediate variables).

Luk

1 Like
3

Many thanks @luk-f-a, yeah indeed it is not mandatory. Good point, at first I thought it was better to use it to initialize the array with a proper shape, however that´s definitelly something I can test the time on.