Итерация по массивам#
Примечание
Массивы поддерживают протокол итератора и могут быть перебраны, как списки Python. См. Индексирование, срезы и итерация раздел в руководстве Quickstart для базового использования и примеров. Остальная часть этого документа представляет nditer объект и охватывает более продвинутое использование.
Объект итератора nditer, представленный в NumPy 1.6, предоставляет множество гибких способов систематического обхода всех элементов одного или нескольких массивов. Эта страница знакомит с некоторыми базовыми способами использования объекта для вычислений над массивами в Python, а затем завершается объяснением, как можно ускорить внутренний цикл в Cython. Поскольку Python-интерфейс
nditer является относительно простым отображением API итератора C-массива, эти идеи также помогут работать с итерацией массива из 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='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’ или 'только для записи'
флаги на операнд.
Затем nditer будет выдавать перезаписываемые буферные массивы, которые вы можете изменять. Однако, поскольку nditer должен скопировать эти данные буфера обратно в исходный массив после завершения итерации, вы должны сигнализировать о завершении итерации одним из двух методов. Вы можете либо:
использовал nditer как контекстный менеджер с помощью
withоператор, и временные данные будут записаны обратно при выходе из контекста.вызвать итератор
closeметод после завершения итерации, который вызовет запись обратно.
nditer больше нельзя итерировать, как только close вызывается или его
контекст завершается.
Пример
>>> 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 могут использоваться на больших блоках посещаемых элементов.
The 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 "" , line 1, in Альтернативный цикл и доступ к элементам#
Чтобы сделать его свойства более доступными во время итерации,
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]
Итерация как конкретный тип данных#
Бывают случаи, когда необходимо обрабатывать массив как другой тип данных, отличный от того, в котором он хранится. Например, может потребоваться выполнять все вычисления с 64-битными числами с плавающей запятой, даже если манипулируемые массивы являются 32-битными числами с плавающей запятой. За исключением написания низкоуровневого кода на C, обычно лучше позволить итератору обрабатывать копирование или буферизацию, а не самостоятельно приводить тип данных во внутреннем цикле.
Существует два механизма, позволяющих это сделать: временные копии и режим буферизации. С временными копиями создаётся копия всего массива с новым типом данных, затем итерация выполняется в копии. Запись разрешена через режим, который обновляет исходный массив после завершения всей итерации. Основной недостаток временных копий заключается в том, что временная копия может потреблять большой объём памяти, особенно если тип данных итерации имеет больший размер элемента, чем исходный.
Режим буферизации смягчает проблему использования памяти и более дружелюбен к кэшу, чем создание временных копий. За исключением особых случаев, когда весь массив нужен сразу вне итератора, буферизация рекомендуется вместо временного копирования. В NumPy буферизация используется ufuncs и другими функциями для поддержки гибких входных данных с минимальными накладными расходами памяти.
В наших примерах мы будем обрабатывать входной массив с комплексным типом данных, чтобы можно было извлекать квадратные корни из отрицательных чисел. Без включения копий или режима буферизации итератор вызовет исключение, если тип данных не совпадает точно.
Пример
>>> 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 "" , line 1, in В режиме копирования '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 для определения, разрешено ли конкретное преобразование. По умолчанию он применяет 'безопасное' приведение. Это означает, например, что он вызовет исключение, если попытаться обработать массив 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 "" , line 1, in >>> 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 "" , line 1, in Одна вещь, на которую стоит обратить внимание, — это преобразования обратно в исходный тип данных при использовании операнда для чтения-записи или только записи. Частый случай — реализация внутреннего цикла с использованием 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 "" , line 2, in Итерация по массиву с трансляцией#
NumPy имеет набор правил для работы с массивами, имеющими различные формы, которые применяются всякий раз, когда функции принимают несколько операндов, комбинируемых поэлементно. Это называется
вещание. The 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 использует флаги 'allocate' и 'writeonly' для операндов, передаваемых как None. Это означает, что мы смогли предоставить только два операнда итератору, и он обработал остальное.
При добавлении параметра '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должны быть включены в старые категории. Значения, которые были в
удаленных категориях, будут установлены в NaN 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#
Те, кто хочет получить действительно высокую производительность своих низкоуровневых операций, должны серьёзно рассмотреть прямое использование API итерации, предоставляемого на C, но для тех, кто не знаком с C или C++, Cython является хорошим компромиссом с разумными компромиссами производительности. Для nditer объект, это означает, что итератор берет на себя широковещание, преобразование dtype и буферизацию, в то время как внутренний цикл передается 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',
... '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.])
Чтобы цитонизировать эту функцию, мы заменяем внутренний цикл (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