Python3 / Numpy: ndarray conditional indexing

Question

I'm coming from a MATLAB background and I'm trying to write this in python/numpy:

[l, m, n] = ndgrid(1:size(dct, 1), 1:size(dct, 2), 1:size(dct, 3));
mycell{i, j} = dct(...
    min.^2 <= l.^2 + m.^2 + n.^2 & ...
    l.^2 + m.^2 + n.^2 <= max.^2)';

So what the code is supposed to to is take all the values of the array that have an index (e.g. x,y,z) that have a 2-norm between min and max, i.e. min^2 < x^2 + y^2 + z^2 < max^2

The only thing I could find was about indexing some values of an array with a condition of the value of the array at this index, however I want to index with a condition on the index itself.

I read about broadcasting and the ix_ function and advanced indexing, however i cannot fit the pieces together.

As a good practice, it would be nice to avoid names like `min`, `max` that are already names for builtins. — Divakar, Dec 31 '16 at 15:09

score 5 · Accepted Answer · edited May 23 '17 at 12:01

NumPy offers us to create open meshes, which could replace the 3D meshes with the help of np.ogrid. This would essentially replace the ndgrid part from MATLAB code, but with additional benefits (read on for more details). These meshes could then be squared and added to perform the equivalent of l.^2 + m.^2 + n.^2 without actually creating 3D versions of l, m and n as we had done with ndgrid. This is a huge performance criteria and has been explored in this previous post and it has shown performance benefits.

Thus, porting over to NumPy, we would have -

m,n,r = dct.shape    
x,y,z = np.ogrid[0:m,0:n,0:r]
vals = x**2+y**2+z**2
mycell[i][j] = dct[(min**2 <= vals) & (vals <= max**2)]

Python3 / Numpy: ndarray conditional indexing

1 Answers1