numpy multidimensional indexing and the function 'take'

Question

On odd days of the week I almost understand multidimensional indexing in numpy. Numpy has a function 'take' which seems to do what I want but with the added bonus that I can control what happens if the indexing is out of rangect Specifically, I have a 3-dimensional array to ask as the lookup-table

lut = np.ones([13,13,13],np.bool)

and a 2x2 array of 3-long vectors to act as indexes into the table

arr = np.arange(12).reshape([2,2,3]) % 13 

IIUC, if I were to write lut[arr] then arr is treated as a 2x2x3 array of numbers and when these are used as indexes into lut they each return a 13x13 array. This explains why lut[arr].shape is (2, 2, 3, 13, 13) .

I can make it do what I want by writing

lut[ arr[:,:,0],arr[:,:,1],arr[:,:,2] ] #(is there a better way to write this?)

and now the three terms act as if they have been zipped to produce a 2x2 array of tuples and lut[<tuple>] produces a single element from lut . The final result is a 2x2 array of entries from lut , just what I want.

I have read the documentation for the 'take' function ...

This function does the same thing as “fancy” indexing (indexing arrays using arrays); however, it can be easier to use if you need elements along a given axis.

and

axis : int, optional
The axis over which to select values.

Perhaps naively, I thought that setting axis=2 I would get three values to use as a 3-tuple to perform the lookup but actually

np.take(lut,arr).shape =  (2, 2, 3)
np.take(lut,arr,axis=0).shape =  (2, 2, 3, 13, 13)
np.take(lut,arr,axis=1).shape =  (13, 2, 2, 3, 13)
np.take(lut,arr,axis=2).shape =  (13, 13, 2, 2, 3)

so it's clear I don't understand what is going on. Can anyone show me how to achieve what I want?

Answer 1

We can compute the linear indices and then use np.take -

np.take(lut, np.ravel_multi_index(arr.T, lut.shape)).T

If you are open to alternatives, we can reshape the indices array to 2D , convert to tuples, index into the data array with it, to give us 1D , which could be reshaped back to 2D -

lut[tuple(arr.reshape(-1,arr.shape[-1]).T)].reshape(arr.shape[:2])

Sample run -

In [49]: lut = np.random.randint(11,99,(13,13,13))

In [50]: arr = np.arange(12).reshape([2,2,3])

In [51]: lut[ arr[:,:,0],arr[:,:,1],arr[:,:,2] ] # Original approach
Out[51]: 
array([[41, 21],
       [94, 22]])

In [52]: np.take(lut, np.ravel_multi_index(arr.T, lut.shape)).T
Out[52]: 
array([[41, 21],
       [94, 22]])

In [53]: lut[tuple(arr.reshape(-1,arr.shape[-1]).T)].reshape(arr.shape[:2])
Out[53]: 
array([[41, 21],
       [94, 22]])

We can avoid the double transposing for the np.take approach, like so -

In [55]: np.take(lut, np.ravel_multi_index(arr.transpose(2,0,1), lut.shape))
Out[55]: 
array([[41, 21],
       [94, 22]])

Generalize to multi-dimensional arrays of generic dimensions

This could be generalized to ndarrays of generic no. of dims, like so -

np.take(lut, np.ravel_multi_index(np.rollaxis(arr,-1,0), lut.shape))

The tuple-based approach should work without any change.

Here's a sample run for the same -

In [95]: lut = np.random.randint(11,99,(13,13,13,13))

In [96]: arr = np.random.randint(0,13,(2,3,4,4))

In [97]: lut[ arr[:,:,:,0] , arr[:,:,:,1],arr[:,:,:,2],arr[:,:,:,3] ]
Out[97]: 
array([[[95, 11, 40, 75],
        [38, 82, 11, 38],
        [30, 53, 69, 21]],

       [[61, 74, 33, 94],
        [90, 35, 89, 72],
        [52, 64, 85, 22]]])

In [98]: np.take(lut, np.ravel_multi_index(np.rollaxis(arr,-1,0), lut.shape))
Out[98]: 
array([[[95, 11, 40, 75],
        [38, 82, 11, 38],
        [30, 53, 69, 21]],

       [[61, 74, 33, 94],
        [90, 35, 89, 72],
        [52, 64, 85, 22]]])

Answer 2

The original problem was trying to do a lookup in a table but some of the indexes were out of bounds and I wanted to control the behaviour when this happened.

import numpy as np
lut = np.ones((5,7,11),np.int) # a 3-dimensional lookup table
print("lut.shape = ",lut.shape ) # (5,7,11)

# valid points are in the interior with value 99,
# invalid points are on the faces with value 0
lut[:,:,:] = 0
lut[1:-1,1:-1,1:-1] = 99

# set up an array of indexes with many of them too large or too small
start = -35
arr = np.arange(start,2*11*3+start,1).reshape(2,11,3)

# This solution has the advantage that I can understand what is going on
# and so I can amend it if I need to

# split arr into tuples along axis=2
arrchannels = arr[:,:,0],arr[:,:,1],arr[:,:,2]

# convert into a flat array but clip the values
ravelledarr = np.ravel_multi_index(arrchannels, lut.shape, mode='clip')

# and now turn back into a list of numpy arrays
# (not an array of the original shape )
clippedarr = np.unravel_index( ravelledarr, lut.shape)
print(clippedarr[0].shape,"*",len(clippedarr)) # produces (2, 11) * 3

# and now I can do the lookup with the indexes clipped to fit
print(lut[clippedarr])

# these are more succinct but opaque ways of doing the same
# due to @Divakar and @hjpauli respectively
print( np.take(lut, np.ravel_multi_index(arr.T, lut.shape, mode='clip')).T )
print( lut.flat[np.ravel_multi_index(arr.T, lut.shape, mode='clip')].T )

The actual application was that I had an rgb image containing some grained wood with some markings on it and I had identified a patch of it. I wanted to take the set of pixels within this patch and mark all points in the whole image that matched one of them. A 256x256x256 existence table is too large so I ran a clustering algorithm on the pixels from the patch and set up existence tables for each cluster (the colours from the patch formed a slender thread through rgb- or hsv-space so the boxes around the clusters were small).

I make the existence tables slightly larger than needed and fill each face with False.

Once I've set up these small existence tables I can now test the rest of the image for matching the patch by looking up each pixel in the tables and use the clipping to make pixels which wouldn't normally map into the tables actually map into a face of the table (and get the value 'False')

Answer 3

I did not try in 3-dimensions. But in 2-dimensions I get the result I want with using numpy.take :

np.take(np.take(T,ix,axis=0), iy,axis=1 )

Perhaps you can expand that to 3-dimensions.

As an example I can adress with two 1-dim arrays for the indices ix and iy the 2-dimensional stencil for the discrete Laplace equation,

ΔT = T[ix-1,iy] + T[ix+1, iy] + T[ix,iy-1] + T[ix,iy+1] - 4*T[ix,iy]

Introducing for a leaner writing:

def q(Φ,kx,ky):
    return np.take(np.take(Φ,kx,axis=0), ky,axis=1 )

then I can run the following python code with numpy.take:

nx = 6; ny= 10
T  = np.arange(nx*ny).reshape(nx, ny)

ix = np.linspace(1,nx-2,nx-2,dtype=int) 
iy = np.linspace(1,ny-2,ny-2,dtype=int)

ΔT = q(T,ix-1,iy)  + q(T,ix+1,iy)  + q(T,ix,iy-1)  + q(T,ix,iy+1)  - 4.0 * q(T,ix,iy)