Python3/Numpy:ndarray条件索引

Question

我来自MATLAB背景,我试图用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)';

所以什么代码应该到是采取具有索引的阵列的所有值(例如X,Y,Z),其具有2范数之间min和max,即min^2 < x^2 + y^2 + z^2 < max^2

我唯一能找到的就是在这个索引上用数组值的条件索引数组的某些值,但是我想索引索引本身的条件.

我读到了有关广播和ix_功能以及高级索引的内容,但是我无法将这些内容整合在一起.

Answer 1

NumPy为我们提供了创建开放网格,可以借助于替换3D网格np.ogrid.这实际上将替换ndgridMATLAB代码中的部分,但具有额外的好处(请继续阅读以获取更多详细信息).然后可以对这些网格进行平方并添加以执行相当的操作,l.^2 + m.^2 + n.^2而无需实际创建3D版本l,m并且n正如我们所做的那样ndgrid.这是一个巨大的性能标准,并且已经进行了探索,this previous post并且已经显示出性能优势.

因此,移植到NumPy,我们会 -

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