Итерация по массивам
Объект-итератор 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 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
© 2005–2019 NumPy Developers
Licensed under the 3-clause BSD License.
https://docs.scipy.org/doc/numpy-1.14.5/reference/arrays.nditer.html
