numpy, different slices for different layers of array

Question

I have 3d array m*n*k and for every 2d-layer I want to take a subarray of size i*j. I have an array c with size 2*k of coordinates of starts of slices for every layer. Is there nice and easy way to get what I need without any loops?

Example:

test = np.arange(18).reshape((3,3,2))
c = np.array([[0,1], [0, 1]])

test[:,:,0] = array([[ 0,  2,  4],
                     [ 6,  8, 10],
                     [12, 14, 16]])

test[:,:,1] = array([[ 1,  3,  5],
                     [ 7,  9, 11],
                     [13, 15, 17]])

I want to get an array

  [[[ 0,  9],
    [ 2,  11]],

   [[ 6,  15],
    [ 8,  17]]]

Solution with loop:

h=2
w=2
layers = 2
F = np.zeros((h,w,layers))
for k in range(layers):
    F[:,:,k] = test[c[0,k]:c[0,k]+h, c[1,k]:c[1,k]+w, k]

Could you share how you would implement that with loops, as it's not clear on getting to the expected o/p from the given inputs? — Divakar
– Divakar, Commented Jan 17, 2017 at 13:43
When I run the code on the given inputs, I see IndexError. — Divakar
– Divakar, Commented Jan 17, 2017 at 14:44

Divakar · Accepted Answer · 2017-01-18 12:48:18Z

1

Here's a vectorized approach making use of broadcasting and advanced-indexing -

d0,d1,d2 = np.ogrid[:h,:w,:layers]
out = test[d0+c[0],d1+c[1],d2]

Sample run -

In [112]: test = np.arange(200).reshape((10,10,2))
     ...: c = np.array([[0,1], [0, 1]])
     ...: 

In [113]: h=4
     ...: w=5
     ...: layers = 2
     ...: F = np.zeros((h,w,layers))
     ...: for k in range(layers):
     ...:     F[:,:,k] = test[c[0,k]:c[0,k]+h, c[1,k]:c[1,k]+w, k]
     ...:     

In [114]: d0,d1,d2 = np.ogrid[:h,:w,:layers]
     ...: out = test[d0+c[0],d1+c[1],d2]
     ...: 

In [115]: np.allclose(F, out)
Out[115]: True

edited Jan 18, 2017 at 12:48

answered Jan 17, 2017 at 15:12

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Egor Ignatenkov Over a year ago

It doesn't seem to work when h != w. I'm getting Index error: shape mismatch

Divakar Over a year ago

@EgorIgnatenkov Was a bug indeed. Updated. Check it out.

Collectives™ on Stack Overflow

numpy, different slices for different layers of array

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related