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

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)