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

Объект итератора 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, end=' ')
...
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, 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-порядка.

Пример

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

• использовать nditer как менеджер контекста с помощью инструкции with, и временные данные будут записаны обратно при выходе из контекста.
• вызвать метод close итератора по завершении итерации, который запустит запись обратно.

nditer больше не может быть перебираемым после того, как был вызван close или после выхода из его контекста.

Пример

>>> 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 будет пытаться предоставить фрагменты, которые как можно больше для внутреннего цикла. Принудительно задавая order=’C’ и ‘F’, мы получим различные размеры внешнего цикла. Этот режим включается путем указания флага итератора.

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

Пример

>>> 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, в зависимости от того, что было запрошено. Примеры ниже показывают выводы, демонстрирующие ход индекса:

Пример

>>> 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 вызовет исключение.

Пример

>>> 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 имеет альтернативную синтаксическую конструкцию для итерации, которая работает явным образом с самим объектом итератора. С этой конструкцией цикла текущее значение доступно с помощью индексации в итератор. Другие свойства, такие как отслеживаемые индексы, остаются прежними. Примеры ниже производят идентичные результаты с предыдущим разделом.

Пример

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

Пример

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

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

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

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

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

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

Пример

>>> 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» задается как флаг для каждого операнда. Это делается для обеспечения контроля на уровне каждого операнда. Режим буферизации задается как флаг итератора.

Пример

>>> 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, чтобы определить, разрешено ли конкретное преобразование. По умолчанию он применяет «безопасное» преобразование. Это означает, например, что он вызовет исключение, если вы попытаетесь обработать массив 64-битных чисел с плавающей точкой как массив 32-битных чисел с плавающей точкой. Во многих случаях правилом «same_kind» наиболее разумно, так как оно позволит преобразование с 64 до 32-битных чисел с плавающей точкой, но не от чисел с плавающей точкой к целым числам или от комплексных чисел к числам с плавающей точкой.

Пример

>>> 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», чтобы разрешить обработку других типов чисел с плавающей точкой. В режиме только для чтения можно предоставить массив целых чисел, а режим чтения-записи вызовет исключение, поскольку преобразование обратно в массив нарушит правило преобразования.

Пример

>>> 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), end=' ')
...
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), end=' ')
...
Traceback (most recent call last):
...
ValueError: operands could not be broadcast together with shapes (2,) (2,3)

Итератор, выделяющий выходные массивы

В функциях NumPy распространен случай, когда выходные данные выделяются на основе вещания входных данных, а также есть необязательный параметр «out», в котором результат будет помещен при его предоставлении. Объект nditer предоставляет удобный способ, который очень упрощает поддержку этой механики.

Покажем, как это работает, создав функцию square, которая возводит свой вход в квадрат. Начнём с минимального определения функции, исключая поддержку параметра «out».

Пример

>>> 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 использует флаги «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']])
...     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 вместо построения другого списка.

Операция во внутреннем цикле — это простое умножение. Всё, что связано с внешним произведением, обрабатывается настройкой итератора.

Пример

>>> 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» предоставлен как флаг итератора.

Для простого примера рассмотрим подсчёт суммы всех элементов в массиве.

Пример

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

Пример

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

Пример

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

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

В качестве примера создадим функцию суммы квадратов. Начнём с реализации этой функции в стандартном Python. Мы хотим поддерживать параметр «ось», аналогичный функции 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',
...                                       '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.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'])
    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 и выражением, использующим встроенную функцию sum 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