Векторизация#

horizontal_line = "_" * 80

Про векторизацию#

NumPy оптимизирован для работы с многомерными массивами, но циклы python нет. В связи с этим распространен подход, называемый векторизацией, при котором устраняются циклы, а вместо них вызываются встроенные в NumPy методы. Таким образом эти циклы осуществляются внутри C кода.

Продемонстрируем это на примере вычисления суммы элементов массива NumPy. Ниже создаётся массив из \(10^7\) случайных чисел.

import numpy as np

N = 10 ** 7

x = np.random.random(N)

Измерим, сколько времени займёт вычисление суммы элементов этого массива в обычном цикле.

Note

%%timeмагическая команда ipython, которая измеряет время, затраченное на исполнение кода в ячейке, в которой она указана.

%%time
s1 = 0
for i in range(N):
    s1 += x[i]
CPU times: total: 3.11 s
Wall time: 3.12 s

Теперь используем для этого метод numpy.sum.

Note

%%timeit — то же самое, что и %%time, но время усредняется по нескольким запускам.

%%timeit
s2 = x.sum()
12 ms ± 75.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Видим значительную разницу в скорости циклов и векторизованного кода.

У векторизации есть три основных плюса:

  1. Скорость: векторизованный код исполняется гораздо быстрее, чем его аналог в циклах;

  2. Количество кода: векторизованный код без циклов как правило короче, а значит в нем сложнее ошибиться;

  3. Выразительность: векторизованный код больше похож на математическое выражение, что упрощает его чтение.

Арифметика.#

Всегда предпочтительнее использовать арифметические операторы сразу к массивам целиком, а не к их элементам внутри циклов.

Заведем два массива длинной \(10^7\) каждый.

x = np.zeros(N)
y = np.ones(N)

Измерим, сколько времени займет вычисление суммы двух этих циклов поэлементно в цикле.

%%time
z = np.zeros(N)
for i in range(N):
    z[i] = x[i] + y[i]
CPU times: total: 4.67 s
Wall time: 4.75 s

Теперь сравним это с векторизованной версией.

%%timeit
z = x + y
43 ms ± 535 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Кроме того, что вычисления заняли гораздо в 103 раза меньше времени, исходный код для сложения выглядит гораздо проще.

Broadcasting#

В ряде случаев допускается применять арифметические операторы к массивам разной формы. Строго говоря, NumPy может оперировать с двумя массивами, если их формы совместимы (broadcastable). Подробнее почитать о broadcasting в документации можно здесь и здесь.

Пример 1. Скаляр и массив.

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

../../_images/broadcasting_1.svg

На картинке выше массив a из трех элементов умножается на скаляр b (который можно представить в виде массива из одного элемента).

a = np.array([1, 2, 3])
b = np.array([2])
result = a * b
print(f"{a} * {b} = {result}")
[1 2 3] * [2] = [2 4 6]

Note

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

Общий случай. В более общем случае два массива совместимы по форме (broadcastable), если размеры вдоль всех их осей совместимы. Размеры вдоль двух осей совместимы, только если

  • они равны;

  • один из них равен единице;

При этом проверка этих размеров происходит справа налево и форма массива меньшей размерности дополняется единицами слева до необходимого количества измерений.

Рассмотрим примеры совместимых и не совместимых форм.

Совместимые формы.

A.shape

B.shape

result.shape

Комментарий

(3,)

(1, )

(3, )

Размер B вдоль единственной оси равен 1, формы A и B совместимы

(1, 3)

(3, 1)

(3, 3)

Размеры вдоль последней оси совместимы, т.к. размер B вдоль нее равен 1. Размеры вдоль первой оси совместимы, т.к. размер A вдоль неё равен 1

(8, 1, 6, 1)

(7, 1, 5)

(8, 7, 5, 6)

Массив B можно представить в виде массива с формой (1, 7, 1, 5). Размеры вдоль всех осей совместимы, т.к. вдоль каждой оси у одного из массивов длинна равняется 1

(15, 4, 3)

(15, 1, 3)

(15, 4, 3)

Размеры вдоль всех осей совместимы

Несовместимые формы.

A.shape

B.shape

Комментарий

(3, )

(4, )

Размеры вдоль единственных осей не равны между собой, ни один из них не равен 1

(2, 1)

(2, 4, 3)

Сначала сравниваются 1 и 3: всё ок. Затем сравниваются 2 и 4: несовместимы. Можно представить форму (2, 1) в виде (1, 2, 1)

Пример 2. Матрица и вектор-строка.

Если складываются (умножаются и т.п.) двумерный и одномерный массивы, то операция возможна, если количество столбцов в матрице соответствует количеству элементов в одномерном массиве.

../../_images/broadcasting_2.svg

В примере выше складываются массивы с формами (4, 3) и (3, ), что эквивалентно записи (4, 3) и (1, 3). Формы совместимы и результат выходит такой, будто одномерный массив расширяется до двумерного массива с формой (4, 3) “по столбцам”.

a = np.array([[0, 0, 0], [10, 10, 10], [20, 20, 20], [30, 30, 30]])
b = np.array([0, 1, 2])
result = a + b

print(f"{a} \n + \n {b} \n = \n {result} ")
[[ 0  0  0]
 [10 10 10]
 [20 20 20]
 [30 30 30]] 
 + 
 [0 1 2] 
 = 
 [[ 0  1  2]
 [10 11 12]
 [20 21 22]
 [30 31 32]] 
../../_images/broadcasting_3.svg

А в этом примере формы не совместимы.

a = np.array([[0, 0, 0], [10, 10, 10], [20, 20, 20], [30, 30, 30]])
b = np.array([1, 2, 3, 4])
a + b
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [14], in <cell line: 3>()
      1 a = np.array([[0, 0, 0], [10, 10, 10], [20, 20, 20], [30, 30, 30]])
      2 b = np.array([1, 2, 3, 4])
----> 3 a + b

ValueError: operands could not be broadcast together with shapes (4,3) (4,) 

Пример 3. Матрица и вектор столбец.

Аналогично можно сложить матрицу с вектором столбцом, но для этого этот массив нужно транспонировать, т.е. привести вектор формы (n, ) (одномерный массив из n элементов) к форме (n, 1) (матрица из n строк длинны 1).

К сожалению, ndarray.T (np.transpose) не приведёт к желаемому результату, если массив одномерный.

print(b, b.T)
[0 1 2] [0 1 2]

Сделать это можно изменив форму или “вставив ось”.

n = b.shape[0]
print(b.reshape((n, 1)))
print()
print(b[:, np.newaxis])
[[1]
 [2]
 [3]
 [4]]

[[1]
 [2]
 [3]
 [4]]

Теперь можно сложить матрицу со столбцом.

result = a + b[:, np.newaxis]
print(result)
[[ 1  1  1]
 [12 12 12]
 [23 23 23]
 [34 34 34]]

Пример 3. Вектор столбец и вектор строка

Если складывать вектор столбец и вектор строку, то каждая из них как бы расширится до матрицы по строкам и по столбцам.

../../_images/broadcasting_4.svg
a = np.array([0, 10, 20, 30])[:, np.newaxis]
b = np.array([1, 2, 3])
result = a + b
print(f"{a} \n + \n {b} \n  = \n {result}")
[[ 0]
 [10]
 [20]
 [30]] 
 + 
 [1 2 3] 
  = 
 [[ 1  2  3]
 [11 12 13]
 [21 22 23]
 [31 32 33]]

Математические функции#

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

\[\begin{split} \mathrm{np.sin}\left( \begin{bmatrix} x_{11} & \ldots & x_{1m} \\ \vdots & \ddots & \vdots \\ x_{n1} & \ldots & x_{nm} \end{bmatrix} \right) = \begin{bmatrix} \mathrm{np.sin}(x_{11}) & \ldots & \mathrm{np.sin}(x_{1m}) \\ \vdots & \ddots & \vdots \\ \mathrm{np.sin}(x_{n1}) & \ldots & \mathrm{np.sin}(x_{nm}) \end{bmatrix} \end{split}\]

Все циклы выполнятся в C коде, а значит быстро.

x = np.random.random(N)
y = np.zeros(N)

Применение функции в цикле поэлементно.

%%time
for i in range(N):
    y[i] = np.sin(x[i])
CPU times: total: 18.5 s
Wall time: 20.2 s

Применение функции сразу ко всему массиву.

%%timeit
y = np.sin(x)
124 ms ± 1.85 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Табуляция функции одной переменной#

Для табуляции функции \(f\colon \mathbb{R} \to \mathbb{R}\) на узлах сетки \(\{x_i, \, i=1,\ldots,N\}\) достаточно просто реализовать функцию \(f\) средствами функций из библиотеки NumPy и передать этой функции массив узлов сетки.

x = np.array([0, np.pi/6, np.pi/2])
y = np.sin(x)

print(y)
[0.  0.5 1. ]

Часто очень полезной в таких ситуациях оказывается полезной метод numpy.linspace, который генерирует равномерную на отрезке сетку. Например, чтобы задать равномерную на отрезке \([-1, 1]\) сетку из 11 точек, достаточно вызывать этот метод в следующей форме.

np.linspace(-1, 1, 11)
array([-1. , -0.8, -0.6, -0.4, -0.2,  0. ,  0.2,  0.4,  0.6,  0.8,  1. ])

Табуляция функции нескольких переменных#

Пусть даны наборы точек \(\{x_i, \, i=1,\ldots,N\}\) и \(\{y_i, \, i=1,\ldots,M\}\) и требуется вычислить значение функции \(f\colon \mathbb{R}^2 \to \mathbb{R}\) в узлах этой сетки:

\[ z_{ji} = f(x_i, y_j), \quad i=1,\ldots,N, \quad j=1,\ldots,M. \]

Не векторизованный код может выглядеть как-то так.

N, M = 2, 3

def f(x, y):
    return np.cos(x) + np.sin(y)

x = np.array([0, 1])
y = np.array([-1, -2, -3])
z = np.zeros((M, N))

for i in range(N):
    for j in range(M):
        z[j, i] = f(x[i], y[j])

print(z)
[[ 0.15852902 -0.30116868]
 [ 0.09070257 -0.36899512]
 [ 0.85887999  0.3991823 ]]

Если функция \(f\) реализована исключительно средствами функций библиотеки NumPy, как в ячейке выше, то можно легко его векторизовать. Для этого удобно применить метод numpy.meshgrid. Принцип его действия в случае двух массивов иллюстрируется формулой

\[\begin{split} \mathrm{np.meshgrid}\left( \begin{bmatrix} x_1 \\ \vdots \\ x_N \end{bmatrix}, \begin{bmatrix} y_1 \\ \vdots \\ y_M \end{bmatrix} \right) = \left( M\left\{ \begin{bmatrix} x_1 & \ldots & x_N \\ \vdots & \ddots & \vdots \\ x_1 & \ldots & x_N \end{bmatrix}\right., \overbrace{ \begin{bmatrix} y_1 & \ldots & y_1 \\ \vdots & \ddots & \vdots \\ y_M & \ldots & y_M \end{bmatrix} }^{\text{N}} \right). \end{split}\]

В более конкретном случае при \(N=2\) и \(M=3\) метод numpy.meshgrid сработает следующим образом.

\[\begin{split} \mathrm{np.meshgrid}\left( \begin{bmatrix} x_1 \\ x_2 \end{bmatrix}, \begin{bmatrix} y_1 \\ y_2 \\ y_3 \end{bmatrix} \right) = \left( \begin{bmatrix} x_1 & x_2 \\ x_1 & x_2 \\ x_1 & x_2 \end{bmatrix}, \begin{bmatrix} y_1 & y_1 \\ y_2 & y_2 \\ y_3 & y_3 \end{bmatrix} \right). \end{split}\]
print(f"{x=}")
print(f"{y=}")

print(horizontal_line)

xv, yv = np.meshgrid(x, y)
print(f"{xv}")
print(f"{yv}")
x=array([0, 1])
y=array([-1, -2, -3])

[[0 1]
 [0 1]
 [0 1]]
[[-1 -1]
 [-2 -2]
 [-3 -3]]

Теперь можно заменить двойной цикл одним вызовом функции.

z = f(xv, yv)
print(z)
[[ 0.15852902 -0.30116868]
 [ 0.09070257 -0.36899512]
 [ 0.85887999  0.3991823 ]]

Агрегирующие функции#

Так же в NumPy реализован ряд агрегирующих методов, таких как сумма элементов массива (ndarray.sum), их произведение (ndarray.prod), минимум (ndarray.min), максимум (ndarray.max) и т.п. Такого рода операции часто применяются в статистике и доступные в NumPy функции такого рода перечисленны по ссылке.

Продемонстрируем поведение таких функций на примере метода numpy.sum.

matrix = np.array([
    [1, 2, 3],
    [10, 20, 30],
    [100, 200, 300]
])
print(matrix)

print(matrix.sum())
[[  1   2   3]
 [ 10  20  30]
 [100 200 300]]
666

По умолчанию все агрегирующие функции игнорируют размерность и форму массива и агрегируют все данные из массива в одно число.

../../_images/sum_1.svg

Каждая из них позволяет не только агрегировать по всему массиву целиком, но и вдоль обозначенной опциональным параметром axis оси. Например, чтобы просуммировать матрицу по столбцам, необходимо передать значение 0 в качестве параметра axis.

print(matrix.sum(axis=0))
[111 222 333]
../../_images/sum_2.svg

А чтобы просуммировать по строкам — значение 1.

print(matrix.sum(axis=1))
[  6  60 600]
../../_images/sum_3.svg