2

I have a 3d tensor and I want to select different slices from the dim=2. something like a[[0, 1], :, [slice(2, 4), slice(1, 3)]].

a=np.arange(2*3*5).reshape(2, 3, 5)
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14]],

       [[15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29]]])
# then I want something like a[[0, 1], :, [slice(2, 4), slice(1, 3)]]
# that gives me np.stack([a[0, :, 2:4], a[1, :, 1:3]]) without a for loop
array([[[ 2,  3],
        [ 7,  8],
        [12, 13]],

       [[16, 17],
        [21, 22],
        [26, 27]]])

and I've seen this and it is not what I want.

Separius
  • 1,226
  • 9
  • 24
  • Why don't you want the linked answers? Do you disagree with my claim that some sort of loop is necessary? Why? – hpaulj Apr 22 '19 at 05:16

2 Answers2

2

You can use advanced indexing as explained here. You will have to pass the row ids which are [0, 1] in your case and the column ids 2, 3 and 1, 2. Here 2,3 means [2:4] and 1, 2 means [1:3]

import numpy as np
a=np.arange(2*3*5).reshape(2, 3, 5)

rows = np.array([[0], [1]], dtype=np.intp)
cols = np.array([[2, 3], [1, 2]], dtype=np.intp)

aa = np.stack(a[rows, :, cols]).swapaxes(1, 2)
# array([[[ 2,  3],
#         [ 7,  8],
#         [12, 13]],

#        [[16, 17],
#         [21, 22],
#         [26, 27]]])

Another equivalent way to avoid swapaxes and getting the result in desired format is 

aa = np.stack(a[rows, :, cols], axis=2).T

A third way I figured out is by passing the list of indices. Here [0, 0] will correspond to [2,3] and [1, 1] will correspond to [1, 2]. The swapaxes is just to get your desired format of output

a[[[0,0], [1,1]], :, [[2,3], [1,2]]].swapaxes(1,2)
Sheldore
  • 37,862
  • 7
  • 57
  • 71
  • 1
    You have converted the slices into advanced indexing arrays - which in the general case will require a loop on the slices and expanding them with `arange` (or `r_`). It should also be mentioned that the `swapaxes` is required because your indexing still has a slice in the middle - and that dimension is 'tacked on' to the end (a case of combined advance and basic indexing). – hpaulj Apr 22 '19 at 05:20
-1

A solution...

import numpy as np
a = np.arange(2*3*5).reshape(2, 3, 5)
np.array([a[0,:,2:4], a[1,:,1:3]])
godot
  • 1,550
  • 16
  • 33