Slicing 3D numpy array using list of index

Question

The objective is to slice 3D array using list of index.

Here, the array is of shape 2,5,5. For simplicity, let assume the index 0 to 4 label as A,B,C,D,E.

Assume we have 3d array as below

array([[[44, 47, 64, 67, 67],
        [ 9, 83, 21, 36, 87],
        [70, 88, 88, 12, 58],
        [65, 39, 87, 46, 88],
        [81, 37, 25, 77, 72]],

       [[ 9, 20, 80, 69, 79],
        [47, 64, 82, 99, 88],
        [49, 29, 19, 19, 14],
        [39, 32, 65,  9, 57],
        [32, 31, 74, 23, 35]]], dtype=int64)

The index of interest is [1,3,4]. Again, we label this as B,D,E`. The expected output, when slicing the 3D array based on the index is as below

array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]], dtype=int64)

However, slicing the array as below

import numpy as np
np.random.seed(0)
arr = np.random.randint(0, 100, size=(2, 5, 5))
k=arr[:,(1,3,4),(1,3,4)]

does not produced the expect output.

In actual use case, the number of element to be sliced is > 3 elements (> B,D,E). Sorry for the lack of correct terminology used

Akshay Sehgal · Accepted Answer · 2021-09-22 19:33:25Z

Try this, which is similar structure to your arr[:,idx,idx] but using np.ix_(). Do read the documentation for np.ix().-

idx = [1,3,4]
ixgrid = np.ix_(idx,idx)
arr[:,ixgrid[0],ixgrid[1]]

array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]])

Explanation

What you are WANT to do is extract a mesh from the last 2 axes of the array. But what you are doing is extract exact indexes from each of the 2 axes.

When you use arr[:,(1,3,4),(1,3,4)], you are essentially asking for (1,1), (3,3) and (4,4) from the two matrices arr[0] and arr[1]

What you need is to extract a mesh. This can be achieved with np.ix_ and the magic of broadcasting.

If you ask for ...

[[1],
 [3],   and   [1,3,4]
 [4]]

... which is what the np.ix_ constructs, you broadcast the indexes and instead ask for a cross product between them, which is (1,1), (1,3), (1,4), (3,1), (3,3)... etc.

Hope that clarifies why you get the result you are getting and how you can actually get what you need.

ddejohn · Accepted Answer · 2021-09-22 18:10:04Z

The problem

Advanced indexing expects all dimensions to be indexed explicitly. What you're doing here is grabbing the elements at coordinates (1, 1), (3, 3), (4, 4) in each array along axis 0.

The solution

What you need to do is this instead:

idx = (1, 3, 4)  # the indices of interest
arr[np.ix_((0, 1), idx, idx)]

Where (0, 1) corresponds to the first two arrays along axis 0.

Output:

array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]], dtype=int64)

As shown above, np.ix_((0, 1), idx, idx)) produces an object which can be used for advanced indexing. The (0, 1) means that you're explicitly selecting the elements from the arrays arr[0] and arr[1]. If you have a more general 3D array of shape (n, m, q) and want to grab the same subarray out of every array along axis 0, you can use

np.ix_(np.arange(arr.shape[0]), idx, idx))

As your indices. Note that idx is repeated here because you wanted those specific indices but in general they don't need to match.

Generalizing

More generally, you can slice and dice however you want like so:

In [1]: arrays_to_select = (0, 1)

In [2]: rows_to_select = (1, 3, 4)

In [3]: cols_to_select = (1, 3, 4)

In [4]: indices = np.ix_(arrays_to_select, rows_to_select, cols_to_select)

In [5]: arr[indices]
Out[5]:
array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]], dtype=int64)

Let's consider some other shape:

In [4]: x = np.random.randint(0, 9, (4, 3, 5))

In [5]: x
Out[5]:
array([[[1, 0, 2, 1, 0],
        [3, 5, 1, 4, 3],
        [1, 8, 1, 4, 2]],

       [[1, 6, 8, 2, 8],
        [0, 0, 4, 2, 3],
        [8, 5, 6, 2, 5]],

       [[4, 4, 8, 6, 0],
        [3, 0, 1, 2, 8],
        [0, 8, 2, 4, 3]],

       [[7, 8, 8, 1, 4],
        [5, 7, 4, 8, 5],
        [7, 5, 5, 3, 4]]])

In [6]: rows = (0, 2)

In [7]: cols = (0, 2, 3, 4)

By using those rows and cols, you'll be grabbing the subarrays composed of all the elements from columns 0 through 4, from only the rows 0 and 2. Let's verify that with the first array along axis 0:

In [8]: arrs = (0,)  # A 1-tuple which will give us only the first array along axis 0

In [9]: x[np.ix_(arrs, rows, cols)]
Out[9]:
array([[[1, 2, 1, 0],
        [1, 1, 4, 2]]])

Now suppose you want the subarrays produced by rows and cols of only the first and last arrays along axis 0. You can explicitly select (0, -1):

In [10]: arrs = (0, -1)

In [11]: x[np.ix_(arrs, rows, cols)]
Out[11]:
array([[[1, 2, 1, 0],
        [1, 1, 4, 2]],

       [[7, 8, 1, 4],
        [7, 5, 3, 4]]])

If, instead, you want that same subarray from all the arrays along axis 0:

In [12]: arrs = np.arange(x.shape[0])

In [13]: arrs
Out[13]: array([0, 1, 2, 3])

In [14]: x[np.ix_(arrs, rows, cols)]
Out[14]:
array([[[1, 2, 1, 0],
        [1, 1, 4, 2]],

       [[1, 8, 2, 8],
        [8, 6, 2, 5]],

       [[4, 8, 6, 0],
        [0, 2, 4, 3]],

       [[7, 8, 1, 4],
        [7, 5, 3, 4]]])

Collectives™ on Stack Overflow

Slicing 3D numpy array using list of index

2 Answers 2

Explanation

Comments

The problem

The solution

Generalizing

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Explanation

Comments

The problem

The solution

Generalizing

Comments

Your Answer

Sign up or log in

Post as a guest

Related