Using np.where to find index of element in 2D array gives ValueError

Question

I'm trying to use np.where to find the index of an element in an array, specifically row number

I have an array of say size 1000 x 6, named 'table'. The first element in each row is a 2 x 2 string array, and the rest are 0s. Eg. a 5 x 6 example of elements in 'table':

    [['s',' ']   0 0 0 0 0
     [' ',' ']]
    [[' ',' ']   0 0 0 0 0
     [' ','a']]
    [[' ',' ']   0 0 0 0 0
     [' ',' ']]         
    [['p',' ']   0 0 0 0 0
     [' ',' ']]
    [[' ',' ']   0 0 0 0 0
     ['b',' ']]  

The 2x2 arrays are all different, and I want to get the index, in particular the row number, of the one containing a specific 2x2 in my large table.

Eg. say I have

    grid = [['s',' ']   
            [' ',' ']]

I would like my code to return [0][0]

I have tried this:

    i,j = np.where(table == grid)

and also

    i,j = np.where(np.all(table == grid))

and i get the following error:

    ValueError: not enough values to unpack (expected 2, got 1)

Using a single value eg.

    index = np.where(table == grid) 

does not result in an error, but print(index) will output an empty array:

    (array([], dtype=int64),)

From similar questions on Stack Overflow I can't seem to figure out how this error applies to mine and I've been staring at it for ages

Any help would be much appreciated

Answer 1

Setup:

b = np.array([['s','t'],['q','r']])
c = np.array([['s',' '],[' ',' ']])
a = np.array([[c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [b,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0]])

Assuming your only interested in column zero; write a function that will test each item in a one-d array. And apply it to column zero

def f(args):
    return [np.all(thing==b) for thing in args]

>>> np.apply_along_axis(f,0,a[:,0])
array([False, False, False, False,  True, False, False, False, False])
>>> 

Use np.where on the result

>>> np.where(np.apply_along_axis(f,0,a[:,0]))
(array([4], dtype=int64),)

Or following the note in the numpy.where docs:

>>> np.asarray(np.apply_along_axis(f,0,a[:,0])).nonzero()
(array([4], dtype=int64),)

---

As @hpaulj points out np.apply_along_axis is not necessary. so

>>> [np.all(thing == b) for thing in a[:,0]]
[False, False, False, False, True, False, False, False, False]

>>> np.asarray([np.all(thing == b) for thing in a[:,0]]).nonzero()
(array([4], dtype=int64),)

And without the Python iteration:

>>> (np.stack(a[:,0])==b).all(axis=(1,2))
array([False, False, False, False,  True, False, False, False, False])

>>> (np.stack(a[:,0])==b).all(axis=(1,2)).nonzero()
(array([4], dtype=int64),)

Answer 2

Here is a solution using vectorize

a = np.array( [np.array([['s',' '],[' ',' ']])  , 0, 0, 0, 0, 0 ])
grid = np.array([['s',' '],[' ',' ']]) 

vfunc = np.vectorize(lambda x: np.all(grid == x))
np.argwhere(vfunc(a))