CODEWITHSUNDEEP
numpysub-array
user3799584
user3799584

Reputation: 927

How to get a subarray in numpy

I have an 3d array and I want to get a sub-array of size (2n+1) centered around an index indx. Using slices I can use

y[slice(indx[0]-n,indx[0]+n+1),slice(indx[1]-n,indx[1]+n+1),slice(indx[2]-n,indx[2]+n+1)]

which will only get uglier if I want a different size for each dimension. Is there a nicer way to do this.

Upvotes: 3

Views: 3834

Answers (2)

cfh
cfh

Reputation: 4666

You don't need to use the slice constructor unless you want to store the slice object for later use. Instead, you can simply do:

y[indx[0]-n:indx[0]+n+1, indx[1]-n:indx[1]+n+1, indx[2]-n:indx[2]+n+1]

If you want to do this without specifying each index separately, you can use list comprehensions:

y[[slice(i-n, i+n+1) for i in indx]]

Upvotes: 2

Divakar
Divakar

Reputation: 221574

You can create numpy arrays for indexing into different dimensions of the 3D array and then use use ix_ function to create indexing map and thus get the sliced output. The benefit with ix_ is that it allows for broadcasted indexing maps. More info on this could be found here. Then, you can specify different window sizes for each dimension for a generic solution. Here's the implementation with sample input data -

import numpy as np

A = np.random.randint(0,9,(17,18,16))  # Input array
indx = np.array([5,10,8])              # Pivot indices for each dim
N = [4,3,2]                            # Window sizes

# Arrays of start & stop indices
start = indx - N
stop = indx + N + 1

# Create indexing arrays for each dimension
xc = np.arange(start[0],stop[0])
yc = np.arange(start[1],stop[1])
zc = np.arange(start[2],stop[2])

# Create mesh from multiple arrays for use as indexing map 
# and thus get desired sliced output
Aout = A[np.ix_(xc,yc,zc)]

Thus, for the given data with window sizes array, N = [4,3,2], the whos info shows -

In [318]: whos
Variable   Type       Data/Info
-------------------------------
A          ndarray    17x18x16: 4896 elems, type `int32`, 19584 bytes
Aout       ndarray    9x7x5: 315 elems, type `int32`, 1260 bytes

The whos info for the output, Aout seems to be coherent with the intended output shape which must be 2N+1.

Upvotes: 2

Related Questions