Filter numpy array according to predicate of position

Question

Suppose I have a two dimensional numpy array with a given shape and I would like to get a view of the values that satisfy a predicate based on the value's position. That is, if x and y are the column and row index accordingly and a predicate x>y the function should return only the array's values for which the column index is greater than the row index.

The easy way to do is a double loop but I would like a possibly faster (vectorized maybe?) approach?

Is there a better way?

Answer 1

In general, you could do this by constructing an open mesh grid corresponding to the row/column indices, apply your predicate to get a boolean mask, then index into your array using this mask:

A = np.zeros((10,20))
y, x  = np.ogrid[:A.shape[0], :A.shape[1]]
mask = x > y
A[mask] = 1

Your specific example happens to be the upper triangle - you can get a copy of it using np.triu, or you can get the corresponding row/column indices using np.triu_indices.