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
– Divakar, Commented Dec 31, 2016 at 15:09

Community · Accepted Answer · 2017-05-23 12:01:54Z

5

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)]

edited May 23, 2017 at 12:01

CommunityBot

11 silver badge

answered Dec 31, 2016 at 14:58

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Python3 / Numpy: ndarray conditional indexing

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related