2

I'm trying to index the last dimension of a 3D matrix with a matrix consisting of indices that I wish to keep.

I have a matrix of thrust values with shape:

(3, 3, 5)

I would like to filter the last index according to some criteria so that it is reduced from size 5 to size 1. I have already found the indices in the last dimension that fit my criteria:

[[0 0 1]
 [0 0 1]
 [1 4 4]]

What I want to achieve: for the first row and first column I want the 0th index of the last dimension. For the first row and third column I want the 1st index of the last dimension. In terms of indices to keep the final matrix will become a (3, 3) 2D matrix like this:

[[0,0,0], [0,1,0], [0,2,1];
 [1,0,0], [1,1,0], [1,2,1];
 [2,0,1], [2,1,4], [2,2,4]]

I'm pretty confident numpy can achieve this, but I'm unable to figure out exactly how. I'd rather not build a construction with nested for loops.

I have already tried:

minValidPower = totalPower[:, :, tuple(indexMatrix)]

But this results in a (3, 3, 3, 3) matrix, so I am not entirely sure how I'm supposed to approach this.

2 Answers2

2

With a as input array and idx as the indexing one -

np.take_along_axis(a,idx[...,None],axis=-1)[...,0]

Alternatively, with open-grids -

I,J = np.ogrid[:idx.shape[0],:idx.shape[1]]
out = a[I,J,idx]
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • `numpy has no attribute 'take_along_axis'` –  Jul 11 '19 at 09:24
  • @LucEvertzen `np.take_along_axis` was introduced in a recent NumPy version. Update NumPy or use [`this solution`](https://stackoverflow.com/a/47048882/) as the replacement or use the second approach. – Divakar Jul 11 '19 at 09:25
0

You can build corresponding index arrays for the first two dimensions. Those would basically be:

[0 1 2]
[0 1 2]
[0 1 2]

[0 0 0] 
[1 1 1] 
[2 2 2]

You can construct these using the meshgrid function. I stored them as m1, and m2 in the example:

vals = np.arange(3*3*5).reshape(3, 3, 5) # test sample
m1, m2 = np.meshgrid(range(3), range(3), indexing='ij')
m3 = np.array([[0, 0, 1], 0, 0, 1], [1, 4, 4]])
sel_vals = vals[m1, m2, m3]

The shape of the result matches the shape of the indexing arrays m1, m2 and m3.

Martin Savc
  • 1
  • 1