python - 使用numpy.nditer()

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

使用numpy.nditer，如何有效地沿着除最后一个以外的所有轴来解析数组x，以便用y（其形状等于x的形状）。 cc>？

以下代码无需使用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循环在速度方面是否会有优势？

最佳答案

您要做的核心是将数组重塑为2d，在一个轴上迭代，然后重塑

_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”，则可以将杂乱的细节隐藏在函数中。 :)

当基础功能可以处理2个维度（一个是主动的，另一个是被动的）时，这种重塑最有用。移调可以根据最方便的方式将活动轴前后移动。

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中，它可以用作以协调的方式迭代几个输入和输出数组的方式，但是它没有任何速度优势。