迭代数组

在NumPy 1.6中引入的迭代器对象nditer提供了许多灵活的方式来以系统化的方式访问一个或多个阵列的所有元素。本页介绍了在Python中使用该对象进行数组计算的基本方法，然后总结了如何加速Cython中的内部循环。由于nditer的Python暴露是C数组迭代器API的相对简单的映射，因此这些想法​​还将帮助处理来自C或C ++的数组迭代。

单数组迭代

使用nditer可以完成的最基本的任务是访问数组的每个元素。使用标准Python迭代器接口逐个提供每个元素。

例

>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a):
...     print x,
...
0 1 2 3 4 5

这个迭代需要注意的一件重要事情是，选择的顺序与数组的内存布局相匹配，而不是使用标准的C或Fortran排序。这是为了提高访问效率，反映了默认情况下，人们只是想访问每个元素而不关心特定的排序。我们可以通过遍历前一个数组的转置来看到这一点，与以C顺序获取该转置的副本相比。

例

>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a.T):
...     print x,
...
0 1 2 3 4 5

>>> for x in np.nditer(a.T.copy(order='C')):
...     print x,
...
0 3 1 4 2 5

a和aT的元素以相同的顺序遍历，即它们存储在内存中的顺序，而aTcopy（order =' C'）以不同的顺序访问，因为它们被放入不同的内存布局中。

控制迭代次序

有时候，按照特定顺序访问数组的元素非常重要，而不管内存中元素的布局如何。nditer对象提供了一个顺序参数来控制迭代的这个方面。具有上述行为的默认值是order ='K'以保持现有订单。对于C命令，可以使用order ='C'和对于Fortran命令的order ='F'来覆盖。

例

>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, order='F'):
...     print x,
...
0 3 1 4 2 5
>>> for x in np.nditer(a.T, order='C'):
...     print x,
...
0 3 1 4 2 5

修改数组值

默认情况下，nditer将输入数组视为只读对象。要修改数组元素，您必须指定读写模式或只写模式。这由每个操作数标志来控制。

Python中的常规赋值只是在本地或全局变量字典中更改引用，而不是修改现有变量。这意味着简单地赋值给x不会将该值放入数组的元素中，而是将x从数组元素引用切换为对该值的引用你分配的。要实际修改数组的元素，x应该使用省略号进行索引。

例

>>> a = np.arange(6).reshape(2,3)
>>> a
array([[0, 1, 2],
       [3, 4, 5]])
>>> for x in np.nditer(a, op_flags=['readwrite']):
...     x[...] = 2 * x
...
>>> a
array([[ 0,  2,  4],
       [ 6,  8, 10]])

使用外部循环

在迄今为止的所有例子中，迭代器一次一个地提供a的元素，因为所有的循环逻辑都在迭代器内部。虽然这很简单方便，但效率不高。更好的方法是将一维最内层循环移动到迭代器外部的代码中。这样，NumPy的矢量化操作就可以用于被访问元素的更大块。

nditer将尝试提供内部循环尽可能大的块。通过强制'C'和'F'的顺序，我们得到不同的外部循环大小。该模式通过指定迭代器标志来启用。

观察到，默认情况下保持本机内存顺序，迭代器能够提供单个一维块，而强制Fortran顺序时，它必须提供三个两个元素块。

例

>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, flags=['external_loop']):
...     print x,
...
[0 1 2 3 4 5]

>>> for x in np.nditer(a, flags=['external_loop'], order='F'):
...     print x,
...
[0 3] [1 4] [2 5]

跟踪索引或多重索引

在迭代过程中，您可能想要在计算中使用当前元素的索引。例如，您可能想按内存顺序访问数组的元素，但使用C顺序，Fortran顺序或多维索引来查找其他数组中的值。

Python迭代器协议没有从迭代器中查询这些附加值的自然方法，因此我们引入了用nditer进行迭代的备用语法。该语法显式地与迭代器对象本身一起工作，所以它的属性在迭代过程中很容易被访问。通过这个循环结构，可以通过索引到迭代器来访问当前值，并且被跟踪的索引是属性index或multi_index，具体取决于请求的内容。

不幸的是，Python交互式解释器在循环的每次迭代期间都会输出while循环内的表达式的值。我们使用这个循环结构修改了示例中的输出，以便更易读。

例

>>> a = np.arange(6).reshape(2,3)
>>> it = np.nditer(a, flags=['f_index'])
>>> while not it.finished:
...     print "%d <%d>" % (it[0], it.index),
...     it.iternext()
...
0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>

>>> it = np.nditer(a, flags=['multi_index'])
>>> while not it.finished:
...     print "%d <%s>" % (it[0], it.multi_index),
...     it.iternext()
...
0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>

>>> it = np.nditer(a, flags=['multi_index'], op_flags=['writeonly'])
>>> while not it.finished:
...     it[0] = it.multi_index[1] - it.multi_index[0]
...     it.iternext()
...
>>> a
array([[ 0,  1,  2],
       [-1,  0,  1]])

跟踪索引或多索引与使用外部循环不兼容，因为它需要每个元素具有不同的索引值。如果试图合并这些标志，nditer对象将引发异常

例

>>> a = np.zeros((2,3))
>>> it = np.nditer(a, flags=['c_index', 'external_loop'])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: Iterator flag EXTERNAL_LOOP cannot be used if an index or multi-index is being tracked

缓冲数组元素

当强制迭代次序时，我们观察到外部循环选项可能以更小的块提供元素，因为元素不能以适当的顺序以恒定的步幅访问。在编写C代码时，这通常很好，但是在纯Python代码中，这会导致性能显着下降。

通过启用缓冲模式，可以使迭代器提供给内部循环的块变大，从而显着减少Python解释器的开销。在强制Fortran迭代顺序的示例中，当启用缓冲时，内部循环可以一次查看所有元素。

例

>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, flags=['external_loop'], order='F'):
...     print x,
...
[0 3] [1 4] [2 5]

>>> for x in np.nditer(a, flags=['external_loop','buffered'], order='F'):
...     print x,
...
[0 3 1 4 2 5]

迭代为特定数据类型

有时需要将数组视为与存储数据类型不同的数据类型。例如，即使被操作的数组是32位浮点数，也可能需要在64位浮点数上进行所有计算。除了编写低级C代码时，通常最好让迭代器处理复制或缓冲，而不是将自己的数据类型转换为内部循环。

有两种机制可以完成这些工作，临时拷贝和​​缓冲模式。对于临时副本，使用新的数据类型创建整个数组的副本，然后在副本中完成迭代。写入访问可以通过在所有迭代完成后更新原始数组的模式来实现。临时副本的主要缺点是临时副本可能会消耗大量内存，特别是如果迭代数据类型的项目大小比原始项目大。

缓冲模式缓解内存使用问题，比制作临时副本更容易缓存。除特殊情况外，整个数组需要在迭代器外部一次使用，建议在临时复制时使用缓冲。在NumPy中，ufuncs和其他函数使用缓冲来支持灵活的输入，并且存储开销最小。

在我们的例子中，我们将用一个复杂的数据类型来处理输入数组，以便我们可以取负数的平方根。没有启用复制或缓冲模式，如果数据类型不准确，迭代器将引发异常。

例

>>> a = np.arange(6).reshape(2,3) - 3
>>> for x in np.nditer(a, op_dtypes=['complex128']):
...     print np.sqrt(x),
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: Iterator operand required copying or buffering, but neither copying nor buffering was enabled

在复制模式下，'copy'被指定为每个操作数标志。这是为了以操作数方式提供控制。缓冲模式被指定为迭代器标志。

例

>>> a = np.arange(6).reshape(2,3) - 3
>>> for x in np.nditer(a, op_flags=['readonly','copy'],
...                 op_dtypes=['complex128']):
...     print np.sqrt(x),
...
1.73205080757j 1.41421356237j 1j 0j (1+0j) (1.41421356237+0j)

>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['complex128']):
...     print np.sqrt(x),
...
1.73205080757j 1.41421356237j 1j 0j (1+0j) (1.41421356237+0j)

迭代器使用NumPy的转换规则来确定是否允许特定的转换。默认情况下，它强制执行“安全”转换。这意味着，例如，如果您尝试将64位浮点数组视为32位浮点数组，则会引发异常。在很多情况下，规则'same_kind'是最合理的规则，因为它允许从64位转换为32位浮点数，但不能从浮点数转换为整数或从复数转换为浮点数。

例

>>> a = np.arange(6.)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32']):
...     print x,
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('float32') according to the rule 'safe'

>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32'],
...                 casting='same_kind'):
...     print x,
...
0.0 1.0 2.0 3.0 4.0 5.0

>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['int32'], casting='same_kind'):
...     print x,
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('int32') according to the rule 'same_kind'

有一点需要注意，当使用读写或只写操作数时，转换回原始数据类型。一个常见的情况是以64位浮点数实现内部循环，并使用'same_kind'强制转换以允许处理其他浮点类型。在只读模式下，可以提供一个整型数组，读写模式将引发异常，因为转换回数组会违反转换规则。

例

>>> a = np.arange(6)
>>> for x in np.nditer(a, flags=['buffered'], op_flags=['readwrite'],
...                 op_dtypes=['float64'], casting='same_kind'):
...     x[...] = x / 2.0
...
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
TypeError: Iterator requested dtype could not be cast from dtype('float64') to dtype('int64'), the operand 0 dtype, according to the rule 'same_kind'

广播数组迭代

NumPy有一套规则来处理具有不同形状的数组，每当函数采用多个操作数组合时，就会应用这些操作数。这被称为broadcasting。当您需要编写这样的函数时，nditer对象可以为您应用这些规则。

作为一个例子，我们打印出一个广播一维和二维数组的结果。

例

>>> a = np.arange(3)
>>> b = np.arange(6).reshape(2,3)
>>> for x, y in np.nditer([a,b]):
...     print "%d:%d" % (x,y),
...
0:0 1:1 2:2 0:3 1:4 2:5

当发生广播错误时，迭代器引发包含输入形状的异常，以帮助诊断问题。

例

>>> a = np.arange(2)
>>> b = np.arange(6).reshape(2,3)
>>> for x, y in np.nditer([a,b]):
...     print "%d:%d" % (x,y),
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: operands could not be broadcast together with shapes (2) (2,3)

迭代器分配的输出数组

NumPy函数中的一个常见情况是根据输入的广播分配输出，另外还有一个可选参数，称为'out'，其结果将在提供时放置。nditer对象提供了一个方便的习惯用法，可以很容易地支持这种机制。

我们将通过创建一个平方形输入的函数square来演示这是如何工作的。让我们从最小的函数定义开始，排除'out'参数支持。

例

>>> def square(a):
...     it = np.nditer([a, None])
...     for x, y in it:
...          y[...] = x*x
...     return it.operands[1]
...
>>> square([1,2,3])
array([1, 4, 9])

默认情况下，nditer将标记'allocate'和'writeonly'用作None传入的操作数。这意味着我们只能向迭代器提供两个操作数，并处理剩余的操作数。

当添加'out'参数时，我们必须显式提供这些标志，因为如果有人将数组作为'out'传入，则迭代器将默认为'readonly'，并且我们的内部循环将失败。'只读'是输入数组的默认原因是为了防止无意中触发减少操作的混淆。如果默认设置为“读写”，则任何广播操作都会触发缩减，本文稍后将讨论该主题。

虽然我们介绍了它，但我们还要介绍'no_broadcast'标志，它会阻止播放的输出。这很重要，因为我们只需要每个输出的一个输入值。汇总多个输入值是需要特殊处理的减少操作。它会引发错误，因为必须在迭代器标志中明确地启用减少，但对于最终用户而言，禁用广播导致的错误消息更容易理解。要了解如何将平方函数推广到约简，请查看关于Cython的部分中的平方和函数。

为了完整性，我们还将添加'external_loop'和'buffered'标志，因为这些是您通常因性能原因需要的标志。

例

>>> def square(a, out=None):
...     it = np.nditer([a, out],
...             flags = ['external_loop', 'buffered'],
...             op_flags = [['readonly'],
...                         ['writeonly', 'allocate', 'no_broadcast']])
...     for x, y in it:
...         y[...] = x*x
...     return it.operands[1]
...

>>> square([1,2,3])
array([1, 4, 9])

>>> b = np.zeros((3,))
>>> square([1,2,3], out=b)
array([ 1.,  4.,  9.])
>>> b
array([ 1.,  4.,  9.])

>>> square(np.arange(6).reshape(2,3), out=b)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 4, in square
ValueError: non-broadcastable output operand with shape (3) doesn't match the broadcast shape (2,3)

外部产品迭代

任何二进制操作都可以像outer中的外部产品一样扩展到数组操作，并且nditer对象提供了一种方法来通过显式映射操作数。也可以使用newaxis索引来做到这一点，但我们将向您展示如何直接使用nditer op_axes参数来完成此操作，而不需要中间视图。

我们将做一个简单的外部产品，将第一个操作数的维度放在第二个操作数的维度之前。The op_axes parameter needs one list of axes for each operand, and provides a mapping from the iterator’s axes to the axes of the operand.

假设第一个操作数是一维的，第二个操作数是二维的。迭代器将有三个维度，所以op_axes将包含两个3元素列表。第一个列表选取第一个操作数的一个轴，对于其余的迭代器轴为-1，最终结果为[0，-1，-1]。第二个列表选取第二个操作数的两个轴，但不应该与第一个操作数中选取的轴重叠。它的列表是[-1，0，1]。输出操作数以标准方式映射到迭代器轴上，因此我们可以提供None而不是构建另一个列表。

内循环中的操作是一个直接的乘法。与外部产品有关的所有事情都由迭代器设置来处理。

例

>>> a = np.arange(3)
>>> b = np.arange(8).reshape(2,4)
>>> it = np.nditer([a, b, None], flags=['external_loop'],
...             op_axes=[[0, -1, -1], [-1, 0, 1], None])
>>> for x, y, z in it:
...     z[...] = x*y
...
>>> it.operands[2]
array([[[ 0,  0,  0,  0],
        [ 0,  0,  0,  0]],
       [[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],
       [[ 0,  2,  4,  6],
        [ 8, 10, 12, 14]]])

减少迭代

只要可写操作数的元素少于整个迭代空间，该操作数正在进行缩减。nditer对象要求任何缩减操作数都被标记为读写，并且只有当'reduce_ok'作为迭代器标志被提供时才允许减少。

举个简单的例子，考虑一下数组中所有元素的总和。

例

>>> a = np.arange(24).reshape(2,3,4)
>>> b = np.array(0)
>>> for x, y in np.nditer([a, b], flags=['reduce_ok', 'external_loop'],
...                     op_flags=[['readonly'], ['readwrite']]):
...     y[...] += x
...
>>> b
array(276)
>>> np.sum(a)
276

组合约简和分配操作数时，情况会有点棘手。在迭代开始之前，任何缩减操作数都必须初始化为其初始值。以下是我们如何做到这一点，并沿a的最后一个轴求和。

例

>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok', 'external_loop'],
...             op_flags=[['readonly'], ['readwrite', 'allocate']],
...             op_axes=[None, [0,1,-1]])
>>> it.operands[1][...] = 0
>>> for x, y in it:
...     y[...] += x
...
>>> it.operands[1]
array([[ 6, 22, 38],
       [54, 70, 86]])
>>> np.sum(a, axis=2)
array([[ 6, 22, 38],
       [54, 70, 86]])

要进行缓冲缩减，需要在安装过程中进行另一次调整。通常情况下，迭代器的构造包括将数据的第一个缓冲区从可读数组复制到缓冲区。任何约简操作数都是可读的，所以它可能被读入缓冲区。不幸的是，该缓冲操作完成后操作数的初始化将不会反映在迭代开始的缓冲区中，并且会产生垃圾结果。

迭代器标志“delay_bufalloc”允许迭代器分配的约简操作数与缓冲一起存在。当这个标志被设置时，迭代器会将其缓冲区保持未初始化状态，直到它接收到一个复位为止，之后它将准备好进行常规迭代。以前的例子看起来是如何启用缓冲的。

例

>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok', 'external_loop',
...                                  'buffered', 'delay_bufalloc'],
...             op_flags=[['readonly'], ['readwrite', 'allocate']],
...             op_axes=[None, [0,1,-1]])
>>> it.operands[1][...] = 0
>>> it.reset()
>>> for x, y in it:
...     y[...] += x
...
>>> it.operands[1]
array([[ 6, 22, 38],
       [54, 70, 86]])

将内环放入Cython

那些希望从低级操作中获得非常好的性能的人应该直接使用C中提供的迭代API，但对于那些不熟悉C或C ++的人来说，Cython是一个很好的中间地带，具有合理的性能折衷。对于nditer对象，这意味着让迭代器负责广播，dtype转换和缓冲，同时给Cython内循环。

对于我们的例子，我们将创建一个平方和函数。首先，让我们直接用Python实现这个函数。我们希望支持类似于numpy sum函数的'axis'参数，所以我们需要为op_axes参数构造一个列表。这看起来如何。

例

>>> def axis_to_axeslist(axis, ndim):
...     if axis is None:
...         return [-1] * ndim
...     else:
...         if type(axis) is not tuple:
...             axis = (axis,)
...         axeslist = [1] * ndim
...         for i in axis:
...             axeslist[i] = -1
...         ax = 0
...         for i in range(ndim):
...             if axeslist[i] != -1:
...                 axeslist[i] = ax
...                 ax += 1
...         return axeslist
...
>>> def sum_squares_py(arr, axis=None, out=None):
...     axeslist = axis_to_axeslist(axis, arr.ndim)
...     it = np.nditer([arr, out], flags=['reduce_ok', 'external_loop',
...                                       'buffered', 'delay_bufalloc'],
...                 op_flags=[['readonly'], ['readwrite', 'allocate']],
...                 op_axes=[None, axeslist],
...                 op_dtypes=['float64', 'float64'])
...     it.operands[1][...] = 0
...     it.reset()
...     for x, y in it:
...         y[...] += x*x
...     return it.operands[1]
...
>>> a = np.arange(6).reshape(2,3)
>>> sum_squares_py(a)
array(55.0)
>>> sum_squares_py(a, axis=-1)
array([  5.,  50.])

为了使用Cython这个函数，我们用专门用于float64 dtype的Cython代码替换了内部循环（y [...] + = x * x）。启用'external_loop'标志后，提供给内部循环的数组将始终为一维，因此需要进行很少的检查。

以下是sum_squares.pyx的列表：

import numpy as np
cimport numpy as np
cimport cython

def axis_to_axeslist(axis, ndim):
    if axis is None:
        return [-1] * ndim
    else:
        if type(axis) is not tuple:
            axis = (axis,)
        axeslist = [1] * ndim
        for i in axis:
            axeslist[i] = -1
        ax = 0
        for i in range(ndim):
            if axeslist[i] != -1:
                axeslist[i] = ax
                ax += 1
        return axeslist

@cython.boundscheck(False)
def sum_squares_cy(arr, axis=None, out=None):
    cdef np.ndarray[double] x
    cdef np.ndarray[double] y
    cdef int size
    cdef double value

    axeslist = axis_to_axeslist(axis, arr.ndim)
    it = np.nditer([arr, out], flags=['reduce_ok', 'external_loop',
                                      'buffered', 'delay_bufalloc'],
                op_flags=[['readonly'], ['readwrite', 'allocate']],
                op_axes=[None, axeslist],
                op_dtypes=['float64', 'float64'])
    it.operands[1][...] = 0
    it.reset()
    for xarr, yarr in it:
        x = xarr
        y = yarr
        size = x.shape[0]
        for i in range(size):
           value = x[i]
           y[i] = y[i] + value * value
    return it.operands[1]

在这台机器上，将.pyx文件构建成模块如下所示，但您可能需要找到一些Cython教程来告诉您系统配置的细节。:

$ cython sum_squares.pyx
$ gcc -shared -pthread -fPIC -fwrapv -O2 -Wall -I/usr/include/python2.7 -fno-strict-aliasing -o sum_squares.so sum_squares.c

从Python解释器运行此解决方案会产生与我们的本机Python / NumPy代码相同的答案。

例

>>> from sum_squares import sum_squares_cy
>>> a = np.arange(6).reshape(2,3)
>>> sum_squares_cy(a)
array(55.0)
>>> sum_squares_cy(a, axis=-1)
array([  5.,  50.])

在IPython中做一点时间表明，减少Cython内部循环的开销和内存分配，对于直接的Python代码和使用NumPy的内置求和函数的表达式提供了非常好的加速。:

>>> a = np.random.rand(1000,1000)

>>> timeit sum_squares_py(a, axis=-1)
10 loops, best of 3: 37.1 ms per loop

>>> timeit np.sum(a*a, axis=-1)
10 loops, best of 3: 20.9 ms per loop

>>> timeit sum_squares_cy(a, axis=-1)
100 loops, best of 3: 11.8 ms per loop

>>> np.all(sum_squares_cy(a, axis=-1) == np.sum(a*a, axis=-1))
True

>>> np.all(sum_squares_py(a, axis=-1) == np.sum(a*a, axis=-1))
True