Spec-Zone.ru › NumPy 2.0

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

Примечание

Массивы поддерживают протокол итератора и могут быть перебираемы как списки Python. Обратитесь к разделу Индексирование, срезы и итерация в руководстве по быстрому началу работы для основного использования и примеров. Остальная часть этого документа представляет собой объект nditer и рассматривает более продвинутое использование.

Объект итератора 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 будет пытаться предоставить фрагменты наибольшего размера для внутреннего цикла. Принудительное использование порядков ‘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. Мы хотим поддерживать параметр «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.])

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

© 2005–2024 NumPy Developers
Licensed under the 3-clause BSD License.
https://numpy.org/doc/2.0/reference/arrays.nditer.html

Spec-Zone.ru

Настройки Оффлайн Что нового Помощь О нас
Spec-Zone .ru
спецификации, руководства, описания, API