Итерация по массивам

Объект-итератор nditer, представленный в NumPy 1.6, предоставляет множество гибких способов систематического посещения всех элементов одного или нескольких массивов. Эта страница представляет некоторые основные способы использования объекта для вычислений над массивами в Python, а затем завершается тем, как можно ускорить внутренний цикл на Cython. Поскольку Python-интерфейс nditer представляет собой относительно прямое отображение API C-итератора массивов, эти идеи также помогут при работе с итерацией по массивам из 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 , так и a.T проходят итерацию в одном и том же порядке, а именно в порядке их хранения в памяти, тогда как элементы a.T.copy(order=’C’) посещаются в другом порядке, потому что они были размещены в другой структуре памяти.

Управление порядком итерации

Бывают случаи, когда важно посещать элементы массива в определенном порядке, независимо от расположения элементов в памяти. Объект nditer предоставляет параметр order для управления этим аспектом итерации. По умолчанию, с поведением, описанным выше, используется order=’K’ для сохранения существующего порядка. Это можно переопределить с order=’C’ для C-порядка и order=’F’ для Fortran-порядка.

Пример

>>> 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]

Итерация как определенного типа данных

Бывают случаи, когда необходимо рассматривать массив как тип данных, отличный от того, в котором он хранится. Например, можно выполнить все вычисления с плавающей точкой двойной точности (64-бит), даже если обрабатываемые массивы представляют собой числа с плавающей точкой одинарной точности (32-бит). За исключением написания кода на низком уровне на C, в большинстве случаев лучше позволить итератору обрабатывать копирование или буферизацию вместо преобразования типа данных в внутреннем цикле.

Существует два механизма для этого: временные копии и буферизация. При использовании временных копий создается копия всего массива с новым типом данных, а затем итерация выполняется по копии. Доступ для записи разрешен через режим, который обновляет исходный массив после завершения всей итерации. Основным недостатком временных копий является то, что временная копия может потреблять большое количество памяти, особенно если тип данных итерации имеет больший размер элемента, чем исходный.

Режим буферизации устраняет проблему с использованием памяти и более дружественен к кэшу, чем создание временных копий. За исключением специальных случаев, когда весь массив нужен сразу за пределами итератора, рекомендуется использовать буферизацию вместо создания временных копий. В NumPy буферизация используется функциями ufunc и другими функциями для поддержки гибких входных данных с минимальным расходом памяти.

В наших примерах мы будем обрабатывать входной массив с комплексным типом данных, чтобы мы могли извлекать квадратные корни из отрицательных чисел. Без включения режимов копирования или буферизации итератор вызовет исключение, если тип данных не совпадает точно.

Пример

>>> 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 имеет набор правил для обработки массивов с различными формами, которые применяются всякий раз, когда функции принимают несколько операндов, которые комбинируются поэлементно. Это называется широковещательной передачей. Объект 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’, и наш внутренний цикл завершится с ошибкой. Причиной, по которой ‘readonly’ является значением по умолчанию для входных массивов, является предотвращение случайного запуска операции сокращения. Если значением по умолчанию было бы ‘readwrite’, любая операция расширения также запускала бы сокращение, тема, которая рассматривается позже в этом документе.

Пока мы тут, давайте также представим флаг ‘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 для достижения этого без промежуточных представлений.

Мы выполним простое внешнее произведение, разместив измерения первого операнда перед измерениями второго операнда. Параметр op_axes требует одного списка осей для каждого операнда и предоставляет отображение осей итератора на оси операнда.

Предположим, что первый операнд одномерный, а второй — двумерный. Итератор будет иметь три измерения, поэтому 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

Те, кто хочет действительно высокой производительности своих низкоуровневых операций, должны серьёзно рассмотреть прямой использование API итераций, предоставленного на C, но для тех, кто не знаком с C или C++, Cython является хорошим средством с приемлемыми компромиссами в производительности. Для объекта nditer это означает предоставление итератору заботы о расширении, преобразовании типов и буферизации, в то время как внутренний цикл передаётся Cython.

Для нашего примера мы создадим функцию суммы квадратов. Начнём с реализации этой функции в обычном Python. Мы хотим поддерживать параметр ‘axis’, подобный функции numpy sum, поэтому нам понадобится создать список для параметра 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, мы заменяем внутренний цикл (y[...] += x*x) кодом Cython, специализированным для типа данных float64. С включенным флагом ‘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