[Python] Apply 1d-mask on numpy 3d-array

I have the following 3d-numpy.ndarray:

import numpy as np
X = np.array([
[[0.0, 0.4, 0.6, 0.0, 0.0],
[0.6, 0.0, 0.0, 0.0, 0.0],
[0.4, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.6, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.4, 0.0, 1.0]],
[[0.1, 0.5, 0.4, 0.0, 0.0],
[0.6, 0.0, 0.0, 0.0, 0.0],
[0.2, 0.0, 0.0, 0.0, 0.0],
[0.1, 0.6, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.4, 0.0, 1.0]]
])


I want a new array where all the rows and columns are dropped where the diagnonal is equal to 1.

idx = np.diag(X[0]) == 1 # for my implementation it is sufficient to look at X[0]


Important to note is that X.shape[1] == X.shape[2], so I try to use the mask as follows

Y = X[:, ~idx, ~idx]


The above returns something different than my desired output:

[[0.0, 0.4, 0.6],
[0.6, 0.0, 0.0],
[0.4, 0.0, 0.0]],
[[0.1, 0.5, 0.4],
[0.6, 0.0, 0.0],
[0.2, 0.0, 0.0]]​

Please advice

Continue reading...