0

Suppose I have a 3D array (tensor)

print a.shape
(100, 100, 100)

and want to index its first axis using one array:

print i.shape
(20,)

and its last axis using another array

print j.shape
(30,)

My intent is to get a (20, 100, 30) array, for example, to be used in assignments like

a[?!, :, ?!] = b

However, I can't figure out how.

print a[i.reshape(20, 1, 1), :, j.reshape(1, 1, 30)].shape
(20, 1, 30, 100)

print a[i.reshape(20, 1, 1), :, j].shape
(20, 1, 30, 100)

print a[i.reshape(20, 1), :, j].shape
(20, 30, 100)

print a[i.reshape(20, 1, 1), j.reshape(1, 1, 30)].shape
(20, 1, 30, 100)

As I understand the "advanced indexing" rules, the first attempt should have worked, but I didn't even end up with a 3D array, and the full dimension (100) came at the end instead of the middle.

MWB
  • 11,740
  • 6
  • 46
  • 91

1 Answers1

1

Approach #1 : You can use np.ix_ for getting such meshes of indices and simply indexing into the input array must give you the desired output. Thus, an implementation would be like so -

a[np.ix_(i,np.arange(a.shape[1]),j)]

Approach #2 : Simpler way using broadcasted indexing -

a[i[:,None],:,j].swapaxes(1,2)

Verify with a sample case -

In [24]: a = np.random.randint(0,99,(5,3,5))

In [25]: i = np.random.randint(0,5,(2))

In [26]: j = np.random.randint(0,5,(2))

In [27]: a[i[0],:,j[0]]
Out[27]: array([15,  7, 74])

In [28]: a[i[0],:,j[1]]
Out[28]: array([32, 19, 85])

In [29]: a[i[1],:,j[0]]
Out[29]: array([76, 65, 96])

In [30]: a[i[1],:,j[1]]
Out[30]: array([54, 65, 66])

In [31]: a[np.ix_(i,np.arange(a.shape[1]),j)]
Out[31]: 
array([[[15, 32],
        [ 7, 19],
        [74, 85]],

       [[76, 54],
        [65, 65],
        [96, 66]]])

In [50]: a[i[:,None],:,j].swapaxes(1,2)
Out[50]: 
array([[[15, 32],
        [ 7, 19],
        [74, 85]],

       [[76, 54],
        [65, 65],
        [96, 66]]])

Assigning values with the indexing

For approach #1, it's just straight-forward -

a[np.ix_(i,np.arange(a.shape[1]),j)] = b

For approach #2, if b is a scalar, it should be straight-forward too -

a[i[:,None],:,j] = b

For approach #2 again, if you are assigning to a ndarray b of shape (20,100,30), we need to swap axes of b before assigning, like so -

a[i[:,None],:,j] = np.swapaxes(b,1,2)
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Approach 2 seems wrong: a[i[:,None],:,j].shape == (20, 30, 100) – MWB Aug 06 '16 at 08:55
  • @MaxB Yup, needed a permute of dims. Updated code. – Divakar Aug 06 '16 at 08:57
  • "swapaxes" can't be used on the left hand side of an assignment – MWB Aug 06 '16 at 08:59
  • @MaxB Could be an old version of NumPy causing that error. Try : `np.swapaxes(a[i[:,None],:,j],1,2)` instead. – Divakar Aug 06 '16 at 09:01
  • Could be my Python version. The error is "SyntaxError: can't assign to function call". Approach 1 works though, it's just a little unsatisfying that one needs to create a temporary array. – MWB Aug 06 '16 at 09:10
  • @MaxB See if the added section on assigning works out for you. – Divakar Aug 06 '16 at 09:22
  • It's strange that `swapaxes` is needed here. Is this a bug/misfeature in numpy? – MWB Aug 06 '16 at 09:44
  • @MaxB Agreed that it's not very intuitive, but from what I gathered, it seems all axes that are not "fancy-indexed", the ones with `colons :` are pushed back. Therefore, `np.ix_` makes more intuitive sense. One hybrid approach that could make more sense would be like : `a[i[:,None,None],np.arange(a.shape[1])[:,None],j] = b` as we are fancy-indexing all axes in this case. But yeah I don't have a good explanation on the seen behavior with a mix of colons and fancy indexing. – Divakar Aug 06 '16 at 10:02