Question

我有一个形状为x的numpy.ndarray (...,3)，即具有任意数量的轴，最后一个轴的已知大小为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＆＃39;对我来说。

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

Answer 1

你正在做的核心是将数组重塑为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)

基本上是相同的策略。如果它看起来不够充分＆＃39; pythonic＆＃39;你可以在一个函数中隐藏凌乱的细节。 :)

当基础功能可以处理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中，它可以作为一种以协调方式迭代多个输入和输出数组的方式，但它没有任何速度优势。

使用numpy.nditer（）

1 个答案: