迭代数组#
备注
数组支持迭代器协议,可以像 Python 列表一样进行迭代.有关基本用法和示例,请参见快速入门指南中的 索引,切片和迭代 部分.本文档的其余部分介绍了 nditer 对象,并涵盖了更高级的用法.
NumPy 1.6 中引入的迭代器对象 nditer 提供了许多灵活的方式,可以系统地访问一个或多个数组的所有元素.本页介绍了一些使用该对象在 Python 中对数组进行计算的基本方法,然后总结了如何在 Cython 中加速内部循环.由于 nditer 的 Python 公开是 C 数组迭代器 API 的相对简单的映射,因此这些想法也将有助于从 C 或 C ++ 进行数组迭代.
单数组迭代#
使用 nditer 可以完成的最基本任务是访问数组的每个元素.使用标准的 Python 迭代器接口逐个提供每个元素.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a):
... print(x, end=' ')
...
0 1 2 3 4 5
对于此迭代,一个需要注意的重要事项是,选择的顺序应与数组的内存布局匹配,而不是使用标准的 C 或 Fortran 排序.这样做是为了提高访问效率,反映了默认情况下人们只想访问每个元素而不必关心特定排序的想法.通过迭代先前数组的转置,与以 C 顺序获取该转置的副本相比,我们可以看到这一点.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a.T):
... print(x, end=' ')
...
0 1 2 3 4 5
>>> for x in np.nditer(a.T.copy(order='C')):
... print(x, end=' ')
...
0 3 1 4 2 5
a 和 a.T 的元素以相同的顺序遍历,即它们存储在内存中的顺序,而 a.T.copy(order=’C’) 的元素以不同的顺序访问,因为它们已被放入不同的内存布局中.
控制迭代顺序#
有时,无论元素在内存中的布局如何,以特定的顺序访问数组的元素都很重要. nditer 对象提供了一个 order 参数来控制迭代的这方面.默认设置(具有上述行为)为 order=’K’,以保持现有顺序.可以使用 order=’C’(对于 C 顺序)和 order=’F’(对于 Fortran 顺序)覆盖此设置.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, order='F'):
... print(x, end=' ')
...
0 3 1 4 2 5
>>> for x in np.nditer(a.T, order='C'):
... print(x, end=' ')
...
0 3 1 4 2 5
修改数组值#
默认情况下, nditer 将输入操作数视为只读对象.为了能够修改数组元素,必须使用 ‘readwrite’ 或 ‘writeonly’ 每操作数标志指定读写或只写模式.
然后,nditer 将产生可写的缓冲区数组,您可以对其进行修改.但是,由于 nditer 必须在迭代完成后将此缓冲区数据复制回原始数组,因此您必须通过以下两种方法之一来指示迭代何时结束.您可以:
一旦调用了 close 或退出了它的上下文,就不能再迭代 nditer.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> a
array([[0, 1, 2],
[3, 4, 5]])
>>> with np.nditer(a, op_flags=['readwrite']) as it:
... for x in it:
... x[...] = 2 * x
...
>>> a
array([[ 0, 2, 4],
[ 6, 8, 10]])
如果您正在编写需要支持旧版本 NumPy 的代码,请注意在 1.15 之前, nditer 不是上下文管理器,也没有 close 方法.相反,它依靠析构函数来启动缓冲区的写回.
使用外部循环#
到目前为止,在所有示例中, a 的元素都是由迭代器一次提供一个,因为所有的循环逻辑都在迭代器内部.虽然这很简单也很方便,但效率不高.更好的方法是将一维最内层循环移到您的代码中,在迭代器外部.这样,NumPy 的向量化操作就可以用于更大块的被访问元素.
nditer 将尝试为内部循环提供尽可能大的块.通过强制 ‘C’ 和 ‘F’ 顺序,我们可以获得不同的外部循环大小.通过指定迭代器标志来启用此模式.
请注意,使用保持本机内存顺序的默认设置,迭代器能够提供单个一维块,而当强制使用 Fortran 顺序时,它必须提供三个块,每个块包含两个元素.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, flags=['external_loop']):
... print(x, end=' ')
...
[0 1 2 3 4 5]
>>> for x in np.nditer(a, flags=['external_loop'], order='F'):
... print(x, end=' ')
...
[0 3] [1 4] [2 5]
跟踪索引或多重索引#
在迭代期间,您可能想要在计算中使用当前元素的索引.例如,您可能想要按内存顺序访问数组的元素,但使用 C 顺序,Fortran 顺序或多维索引来查找不同数组中的值.
索引由迭代器对象本身跟踪,并且可以通过 index 或 multi_index 属性访问,具体取决于请求的内容.下面的示例显示了演示索引进度的打印输出:
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> it = np.nditer(a, flags=['f_index'])
>>> for x in it:
... print("%d <%d>" % (x, it.index), end=' ')
...
0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>
>>> it = np.nditer(a, flags=['multi_index'])
>>> for x in it:
... print("%d <%s>" % (x, it.multi_index), end=' ')
...
0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>
>>> with np.nditer(a, flags=['multi_index'], op_flags=['writeonly']) as it:
... for x in it:
... x[...] = it.multi_index[1] - it.multi_index[0]
...
>>> a
array([[ 0, 1, 2],
[-1, 0, 1]])
跟踪索引或多重索引与使用外部循环不兼容,因为它需要每个元素的不同索引值.如果您尝试组合这些标志,则 nditer 对象将引发异常.
示例
>>> import numpy as np
>>> 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
备用循环和元素访问#
为了使其属性在迭代期间更容易访问, nditer 有一种用于迭代的替代语法,该语法显式地与迭代器对象本身一起使用.使用此循环构造,可以通过索引到迭代器中来访问当前值.其他属性(如跟踪的索引)保持不变.下面的示例产生与上一节中相同的结果.
示例
>>> import numpy as np
>>> 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), end=' ')
... is_not_finished = 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), end=' ')
... is_not_finished = it.iternext()
...
0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>
>>> with np.nditer(a, flags=['multi_index'], op_flags=['writeonly']) as it:
... while not it.finished:
... it[0] = it.multi_index[1] - it.multi_index[0]
... is_not_finished = it.iternext()
...
>>> a
array([[ 0, 1, 2],
[-1, 0, 1]])
缓冲数组元素#
当强制迭代顺序时,我们观察到外部循环选项可能会提供较小的块中的元素,因为不能以恒定的步幅按适当的顺序访问这些元素.在编写 C 代码时,这通常很好,但是在纯 Python 代码中,这会导致性能显着下降.
通过启用缓冲模式,迭代器提供给内部循环的块可以变得更大,从而显着降低了 Python 解释器的开销.在强制 Fortran 迭代顺序的示例中,启用缓冲后,内部循环可以一次看到所有元素.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, flags=['external_loop'], order='F'):
... print(x, end=' ')
...
[0 3] [1 4] [2 5]
>>> for x in np.nditer(a, flags=['external_loop','buffered'], order='F'):
... print(x, end=' ')
...
[0 3 1 4 2 5]
以特定数据类型迭代#
有时需要将数组视为与存储时不同的数据类型.例如,即使被操作的数组是 32 位浮点数,也可能希望对 64 位浮点数进行所有计算.除了编写低级 C 代码外,通常最好让迭代器处理复制或缓冲,而不是在内部循环中自己转换数据类型.
有两种机制可以实现这一点,即临时副本和缓冲模式.使用临时副本,会创建具有新数据类型的整个数组的副本,然后在副本中进行迭代.允许通过一种模式进行写入访问,该模式在所有迭代完成后更新原始数组.临时副本的主要缺点是临时副本可能会消耗大量内存,特别是如果迭代数据类型具有比原始数据类型更大的项目大小.
缓冲模式可以缓解内存使用问题,并且比创建临时副本更利于缓存.除了在迭代器外部一次性需要整个数组的特殊情况外,建议使用缓冲而不是临时复制.在 NumPy 中,ufunc 和其他函数使用缓冲来支持灵活的输入,同时最大程度地减少内存开销.
在我们的示例中,我们将处理具有复杂数据类型的输入数组,以便我们可以对负数进行平方根运算.如果不启用复制或缓冲模式,如果数据类型不完全匹配,迭代器将引发异常.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3) - 3
>>> for x in np.nditer(a, op_dtypes=['complex128']):
... print(np.sqrt(x), end=' ')
...
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’ 被指定为每个操作数的标志.这样做是为了以每个操作数的方式提供控制.缓冲模式被指定为迭代器标志.
示例
>>> import numpy as np
>>> 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), end=' ')
...
1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['complex128']):
... print(np.sqrt(x), end=' ')
...
1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)
迭代器使用 NumPy 的类型转换规则来确定是否允许进行特定的转换.默认情况下,它强制执行 ‘safe’ 转换.例如,这意味着如果您尝试将 64 位浮点数组视为 32 位浮点数组,它将引发异常.在许多情况下,规则 ‘same_kind’ 是最合理的规则,因为它允许从 64 位浮点转换为 32 位浮点,但不允许从浮点转换为整数或从复数转换为浮点.
示例
>>> import numpy as np
>>> a = np.arange(6.)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32']):
... print(x, end=' ')
...
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, end=' ')
...
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, end=' ')
...
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’ 类型转换来允许处理其他浮点类型.虽然在只读模式下,可以提供整数数组,但读写模式会引发异常,因为转换回数组会违反类型转换规则.
示例
>>> import numpy as np
>>> 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 对象可以为您应用这些规则.
作为一个例子,我们打印出一个一维数组和一个二维数组一起广播的结果.
示例
>>> import numpy as np
>>> a = np.arange(3)
>>> b = np.arange(6).reshape(2,3)
>>> for x, y in np.nditer([a,b]):
... print("%d:%d" % (x,y), end=' ')
...
0:0 1:1 2:2 0:3 1:4 2:5
当发生广播错误时,迭代器会引发一个异常,其中包括输入形状以帮助诊断问题.
示例
>>> import numpy as np
>>> a = np.arange(2)
>>> b = np.arange(6).reshape(2,3)
>>> for x, y in np.nditer([a,b]):
... print("%d:%d" % (x,y), end=' ')
...
Traceback (most recent call last):
...
ValueError: operands could not be broadcast together with shapes (2,) (2,3)
迭代器分配的输出数组#
NumPy 函数中一个常见的情况是,根据输入的广播分配输出,并且还有一个名为 ‘out’ 的可选参数,当提供该参数时,结果将放置在该参数中. nditer 对象提供了一个方便的习惯用法,可以非常容易地支持这种机制.
我们将通过创建一个函数 square 来展示它是如何工作的,该函数对其输入进行平方.让我们从一个不包含 ‘out’ 参数支持的最小函数定义开始.
示例
>>> import numpy as np
>>> def square(a):
... with np.nditer([a, None]) as it:
... for x, y in it:
... y[...] = x*x
... return it.operands[1]
...
>>> square([1,2,3])
array([1, 4, 9])
默认情况下, nditer 对作为 None 传入的操作数使用标志 ‘allocate’ 和 ‘writeonly’.这意味着我们只需要向迭代器提供两个操作数,它就会处理其余的事情.
当添加 ‘out’ 参数时,我们必须显式地提供这些标志,因为如果有人将数组作为 ‘out’ 传入,迭代器将默认为 ‘readonly’,并且我们的内部循环将失败.’readonly’ 是输入数组的默认值的原因是为了防止意外触发归约操作的混淆.如果默认值为 ‘readwrite’,任何广播操作也会触发归约,这是一个将在本文档后面介绍的主题.
与此同时,让我们也引入 ‘no_broadcast’ 标志,这将防止输出被广播.这很重要,因为我们只希望每个输出有一个输入值.聚合多个输入值是一种归约操作,需要特殊处理.它已经会引发一个错误,因为必须在迭代器标志中显式启用归约,但是禁用广播导致的错误消息对于最终用户来说更容易理解.要了解如何将平方函数推广到归约,请查看关于 Cython 部分中的平方和函数.
为了完整起见,我们还将添加 ‘external_loop’ 和 ‘buffered’ 标志,因为这些是您通常出于性能原因所需要的.
示例
>>> import numpy as np
>>> def square(a, out=None):
... it = np.nditer([a, out],
... flags = ['external_loop', 'buffered'],
... op_flags = [['readonly'],
... ['writeonly', 'allocate', 'no_broadcast']])
... with it:
... 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):
...
ValueError: non-broadcastable output operand with shape (3,) doesn't
match the broadcast shape (2,3)
外积迭代#
任何二元操作都可以像 outer 中那样,以外积的方式扩展到数组操作,并且 nditer 对象提供了一种通过显式映射操作数的轴来实现此目的的方法.也可以使用 newaxis 索引来实现这一点,但我们将向您展示如何直接使用 nditer 的 op_axes 参数来实现这一点,而无需中间视图.
我们将做一个简单的外积,将第一个操作数的维度放在第二个操作数的维度之前. op_axes 参数需要每个操作数一个轴列表,并提供从迭代器的轴到操作数轴的映射.
假设第一个操作数是一维的,第二个操作数是二维的.迭代器将有三个维度,因此 op_axes 将有两个 3 元素列表.第一个列表选择第一个操作数的一个轴,对于迭代器的其余轴为 -1,最终结果为 [0, -1, -1].第二个列表选择第二个操作数的两个轴,但不应与第一个操作数中选出的轴重叠.它的列表是 [-1, 0, 1].输出操作数以标准方式映射到迭代器轴上,因此我们可以提供 None 而不是构造另一个列表.
内部循环中的操作是一个简单的乘法.与外积有关的所有事情都由迭代器设置处理.
示例
>>> import numpy as np
>>> 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])
>>> with it:
... for x, y, z in it:
... z[...] = x*y
... result = it.operands[2] # same as z
...
>>> result
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]]])
请注意,一旦迭代器关闭,我们就无法访问 operands ,必须使用在上下文管理器内部创建的引用.
归约迭代#
每当可写操作数的元素少于完整迭代空间时,该操作数都将进行归约. nditer 对象要求任何归约操作数都标记为读写,并且仅当提供"reduce_ok"作为迭代器标志时才允许归约.
对于一个简单的例子,考虑取数组中所有元素的和.
示例
>>> import numpy as np
>>> a = np.arange(24).reshape(2,3,4)
>>> b = np.array(0)
>>> with np.nditer([a, b], flags=['reduce_ok'],
... op_flags=[['readonly'], ['readwrite']]) as it:
... for x,y in it:
... y[...] += x
...
>>> b
array(276)
>>> np.sum(a)
276
当组合归约和分配的操作数时,事情会变得有点棘手.在迭代开始之前,任何归约操作数都必须初始化为其起始值.这是我们如何做到这一点的,沿着 a 的最后一个轴求和.
示例
>>> import numpy as np
>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, [0,1,-1]])
>>> with it:
... it.operands[1][...] = 0
... for x, y in it:
... y[...] += x
... result = it.operands[1]
...
>>> result
array([[ 6, 22, 38],
[54, 70, 86]])
>>> np.sum(a, axis=2)
array([[ 6, 22, 38],
[54, 70, 86]])
要进行缓冲归约,需要在设置期间进行另一次调整.通常,迭代器构造涉及将来自可读数组的第一个数据缓冲区复制到缓冲区中.任何归约操作数都是可读的,因此可以将其读入缓冲区.不幸的是,在此缓冲操作完成后对操作数进行初始化将不会反映在迭代开始的缓冲区中,并且会产生垃圾结果.
迭代器标志"delay_bufalloc"用于允许迭代器分配的归约操作数与缓冲一起存在.设置此标志后,迭代器将使其缓冲区保持未初始化状态,直到收到重置为止,之后它将准备好进行常规迭代.如果我们也启用缓冲,则前一个示例的外观如下.
示例
>>> import numpy as np
>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok',
... 'buffered', 'delay_bufalloc'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, [0,1,-1]])
>>> with it:
... it.operands[1][...] = 0
... it.reset()
... for x, y in it:
... y[...] += x
... result = it.operands[1]
...
>>> result
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',
... 'buffered', 'delay_bufalloc'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, axeslist],
... op_dtypes=['float64', 'float64'])
... with it:
... 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.)
>>> sum_squares_py(a, axis=-1)
array([ 5., 50.])
要 Cython 化此函数,我们将内部循环 (y[…] += xx) 替换为专门针对 float64 dtype 定制的 Cython 代码.启用"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'])
with it:
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 内置 sum 函数的表达式更好的加速.
>>> 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