在numpy ndarray中沿指定轴循环向量的有效方法是什么？

Question

我通过沿 numpy ndarray（可以是任何维度）的轴（可以是任何轴）循环向量来处理数据。

我没有直接处理数组，因为数据并不完美。它需要对每个载体进行质量控制。如果不好，向量将被零（或nan）填充，并且没有真正的处理。

我发现this Q 类似，但我的问题要困难得多，因为

1. ndim 是任意的。

对于 3D 数组，我可以像这样沿 axis 1 获取向量

 x = np.arange(24).reshape(2,3,4)
 for i in range(x.shape[0]):
     for k in range(x.shape[2]):
         process(x[i,:,k])

但是如果ndim和取到的axis不固定，怎么取向量呢？

1. 获取向量的轴是任意的。

我正在考虑的一种可能方式是

 y = x.swapaxes(ax,-1)
 # loop over vectors along last axis
 for i in np.ndindex(y.shape[:-1]):
     process(y[i+(slice(None),)])
 # then swap back
 z = y.swapaxes(ax,-1)

但我怀疑这种方法的效率。

Answer 1

测试效率的最佳方法是对实际示例进行时间测试。但是%timeit (ipython) 对玩具示例的测试只是一个开始。

根据回答类似“如果必须重复”问题的经验，时间上没有太大差异。 np.frompyfunc 具有适度的速度优势 - 但它的 pyfunc 采用标量，而不是数组或切片。 （np.vectorize 是这个函数的更好的 API，但速度有点慢）。

但是在这里，您希望将数组的一维切片传递给您的函数，同时遍历所有其他维度。我认为替代迭代方法没有太大区别。

swapaxis、transpose 和 ravel 等操作很快，通常只是创建具有不同形状和步幅的新视图。

np.ndindex 使用np.nditer（带有多索引平面）迭代一系列维度。 nditer 在 C 代码中使用时速度很快，但在 Python 代码中使用时并没有什么特别之处。

np.apply_along_axis 创建一个(i,j,:,k) 索引元组，并对变量进行步进。这是一个很好的通用方法，但没有做任何特别的事情来加快速度。 itertools.product 是另一种生成索引的方式。

但通常不是迭代机制让事情变慢，而是对函数的重复调用。您可以使用普通函数来测试迭代机制，例如

def foo(x):
   return x

====================

您无需swapaxes 即可使用ndindex；您可以使用它在任意轴组合上进行迭代。

例如，制作一个 3d 数组，并沿中间维度求和：

In [495]: x=np.arange(2*3*4).reshape(2,3,4)

In [496]: N=np.ndindex(2,4)

In [497]: [x[i,:,k].sum() for i,k in N]
Out[497]: [12, 15, 18, 21, 48, 51, 54, 57]

In [498]: x.sum(1)
Out[498]: 
array([[12, 15, 18, 21],
       [48, 51, 54, 57]])

我认为这对速度没有影响；代码更简单。

====================

另一个可能的工具是np.ma，屏蔽数组。使用那些您将单个元素标记为已屏蔽（因为它们是 nan 或 0）。它的代码可以评估 sum、mean、product 之类的内容，以使屏蔽值不会损害解决方案。

又是 3d 数组：

In [517]: x=np.arange(2*3*4).reshape(2,3,4)

添加一些不好的值：

In [518]: x[1,1,2]=99    
In [519]: x[0,0,:]=99

这些值搞砸了正常的sum：

In [520]: x.sum(axis=1)
Out[520]: 
array([[111, 113, 115, 117],
       [ 48,  51, 135,  57]])

但如果我们屏蔽它们，它们就会从解决方案中“过滤掉”（在这种情况下，它们被临时设置为 0）

In [521]: xm=np.ma.masked_greater(x,50)

In [522]: xm
Out[522]: 
masked_array(data =
 [[[-- -- -- --]
  [4 5 6 7]
  [8 9 10 11]]

 [[12 13 14 15]
  [16 17 -- 19]
  [20 21 22 23]]],
             mask =
 [[[ True  True  True  True]
 ...
  [False False False False]]],
       fill_value = 999999)

In [523]: xm.sum(1)
Out[523]: 
masked_array(data =
 [[12 14 16 18]
 [48 51 36 57]],
 ...)

Answer 2

你考虑过numpy.nditer吗？

另见Iterating over arrays。

编辑：也许另一种解决方案就是使用：

• flatten
• ravel
• flat 一维迭代器

因此，无论数组的初始暗度如何，您都可以像 1D 一样迭代，然后 reshape 将数组恢复到其原始形状。