Question

我已阅读此post并怀疑我需要更好地理解索引。我想做np.ndenumerate之类的事情，但是（在这种特殊情况下）仅在前三个维度上，返回一系列向量：

x = np.linspace(0., 1., 4)
y = np.linspace(0., 1., 5)
z = np.linspace(0., 1., 2)
X, Y, Z = np.meshgrid(x, y, z)   # define a grid
F = np.zeros(X.shape + (3,))     # a 3D vector field 
F = np.random.rand(5*4*2*3).reshape(5,4,2,3)   # added this later just to give non-zero for testing

thetas = np.linspace(0, 2.*np.pi, 21)[:-1]    # drop the duplicate
v = np.array([(np.cos(theta), np.sin(theta), 0.0) for theta in thetas])

for tup, vec in magical_enumerate(F, axis=(0,1,2)):  # it's magic! (need correct expression)
    F(tup) = np.cross(v, vec).sum(axis=0)     # sum over 20 vectors in v

有没有办法在没有大量循环或列表解释的情况下执行此操作？网格将是大的，所以numpythony和速度是值得赞赏的。那么非顺序尺寸（例如轴=（0,2））？谢谢！

Answer 1

np.ndindex可能会成功。它在一组维度上生成迭代器。

In [231]: F=np.zeros((2,3,3,3))

In [232]: for tup in np.ndindex(F.shape[:3]):
    # vec = data[tup] etc
    F[tup]=tup
   .....:     

In [233]: F
Out[233]: 
array([[[[ 0.,  0.,  0.],
         [ 0.,  0.,  1.],
         [ 0.,  0.,  2.]],

         ...
         [ 1.,  2.,  1.],
         [ 1.,  2.,  2.]]]])

我建议查看ndenumerate和ndindex的代码。 ndindex正在nditer模式下使用较新的multi_index。 ndenumerate使用flat来迭代所有值。

我在其他SO中建议你如何构建自己的multi_index迭代器，模仿ndindex。搜索nditer可能会产生这些内容。

这不会比多个循环提供速度优势，因为你仍在处理相同数量的最内部计算。

对于非顺序维度，同样的ndindex可以使用，但在将其用作索引之前必须操纵tup：

In [243]: for tup in np.ndindex((F.shape[0],F.shape[2])):
    tup1=(tup[0],slice(None),tup[1])
    F[tup]=[tup[0],np.nan,tup[1]]
   .....:

np.apply_along_axis和np.apply_over_axes是在一个或多个轴上生成索引的其他示例，同时切割其他轴。

像partial / selective np.ndenumerate这样的东西？

1 个答案: