python - numpy.nditer() 的使用

Question

我有一个形状为 (...,3) 的 numpy.ndarray x，即具有任意数量的轴，最后一个轴的已知大小为 3 。我还有一个函数 f ，它接受形状 (3) 的数组(实际上是 3D 空间中的一个点)作为参数，并返回另一个形状 的数组(3)(实际上是 3D 空间中的向量)。不幸的是，这个函数不能(至少很容易)被矢量化。

使用numpy.nditer ，如何沿着除最后一个轴之外的所有轴有效解析数组 x ，以便填充数组 y (其形状等于 之一) x)，结果为 f？

以下代码将在不使用 nditer 的情况下完成此操作:

import numpy as np

def f(x):
   '''Simple function for this exemple.
   Can only deal with array of shape (3,)
   '''
   assert x.ndim == 1 and x.shape[0] == 3
   y = np.zeros_like(x)
   y[0] = x[0]
   y[1] = x[1]**2
   y[2] = x[2]**3
   return y

x = np.arange(15).reshape(5,3)

_x = x.reshape(-1,3)
_y = np.zeros_like(_x)
for i in xrange(_x.shape[0]):
    _y[i,:] = f(_x[i,:])
y = _y.reshape(x.shape)

但对我来说看起来并不“Pythonic”。

作为一个额外的问题，使用 nditer 而不是上面的经典 python 循环在速度方面会有优势吗？

Answer 1

您所做的核心是将数组 reshape 为二维，在一个轴上迭代，然后 reshape 回来

_x = x.reshape(-1,3)
_y = np.zeros_like(_x)
for i in xrange(_x.shape[0]):
    _y[i,:] = f(_x[i,:])
y = _y.reshape(x.shape)

将其与 tensordot 的作用进行比较:

newshape_a = (-1, N2)
....
at = a.transpose(newaxes_a).reshape(newshape_a)
bt = b.transpose(newaxes_b).reshape(newshape_b)
res = dot(at, bt)
return res.reshape(olda + oldb)

基本相同的策略。如果它看起来不够“Pythonic”，你可以将困惑的细节隐藏在函数中。 :)

当底层函数可以处理二维(一个是主动的，一个是被动的)“随波逐流”时，这种 reshape 是最有用的。转置可以将事件轴向前或向后移动，具体取决于最方便的方式。

apply_along_axis 中使用的另一个策略是构建索引列表:

for i in range(N):
    fun(arr[tuple([slice(N),slice(N)...,i])]

我过去回答过有关 nditer 的类似问题。 https://stackoverflow.com/a/28727290/901925

np.ndindex 是使用 nditer 迭代轴子集的一个很好的示例。看看它的代码。基本上，它构造了一个正确形状的虚拟数组，并生成 multi_index 类型的索引。

http://docs.scipy.org/doc/numpy/reference/arrays.nditer.html#tracking-an-index-or-multi-index 中也对此进行了说明。

该索引文档是对 nditer 的最佳描述。请注意它如何以 cython 示例结尾。我认为这是 Python nditer 的最佳用途 - 作为在 cython 或 c 中使用它的垫脚石。

在 Python 中，它可以作为一种以协调方式迭代多个输入和输出数组的方法，但它没有任何速度优势。