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.

Improve this question

2 Answers 2

Reset to default
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]
Improve this answer
Sign up to request clarification or add additional context in comments.

2 Comments

numpy has no attribute 'take_along_axis'
@LucEvertzen np.take_along_axis was introduced in a recent NumPy version. Update NumPy or use this solution as the replacement or use the second approach.
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.

Improve this answer

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.