Векторизация на примере вычисления статистик

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

Рассмотрим в качестве примере то же задание, но для вычисления этих же характеристик одного числового ряда. Т.е. будем считать, что есть числовой ряд \(x = \{x_1, ..., x_N\}\), данный нам в виде массива x. Реализуем функцию, которая выдаст нам среднее и среднеквадратичное отклонения этого числового ряда.

Подход в лоб

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

def mean1D(x):
    s = 0 
    for number in x:
        s += number
    return s / x.size

Приведенная выше функция заводит переменную S, в которую на протяжении цикла накапливается сумма. Когда сумма всего массива накоплена, то возвращается результат деления этой суммы на количество элементов в массиве, которое можно узнать у массива x по атрибуту size (предполагаем, что массив одномерный).

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

import math


def std1D(x, mean):
    s = 0
    for number in x:
        s += (number - mean) ** 2
    return math.sqrt(s / x.size)

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

import numpy as np

N = 100 # размер массива
x = np.random.random(N)

mean_1 = mean1D(x)
mean_2 = x.mean()
std_1 = std1D(x, mean_1)
std_2 = x.std()


print(mean_1, mean_2)
print(std_1, std_2)

0.5153244859252695 0.5153244859252698
0.28374330737876574 0.28374330737876574

Векторизованные версии

Функции работают, теперь перепишем их, используя векторизацию.

def mean1D_vectorized(x):
    return x.sum() / x.size

Векторизованная версия функции mean1D примитивна. Цикл суммирования заменили на вызов метода ndarrray.sum.

def std1D_vectorized(x, mean):
    tmp_1 = x - mean
    tmp_2 = tmp_1 ** 2
    tmp_3 = tmp_2.sum()
    return math.sqrt(tmp_3 / x.size)

Вычисление среднеквадратичного отклонения рассмотрим подробнее.

def std1D_vectorized(x, mean):
    tmp_1 = x - mean

Как и прежде, чтобы не вычислять дважды среднее принимаем, mean в качестве второго аргумента. На следующей стоке вычисляется разница между массивом x и средним значением mean (скаляром). Формы не совпадают, а значит производится broadcasting и в результате tmp_1[i] = x[i] - mean.

def std1D_vectorized(x, mean):
    tmp_1 = x - mean
    tmp_2 = tmp_1 ** 2

Далее встречается возведение одномерного массива tmp_1 в степень 2. Т.к. 2 — скаляр, то производится broadcasting и значение tmp_2[i] = tmp_1[i] ** 2 = (x[i] - mean)  ** 2.

def std1D_vectorized(x, mean):
    tmp_1 = x - mean
    tmp_2 = tmp_1 ** 2
    tmp_3 = tmp_2.sum()

В следующей строке вычисляется сумма элементов массива tmp_2, т.е. \(\sum_{i=0}^{N-1}\) tmp_2[i] = \(\sum_{i=0}^{N-1}\)tmp_1[i]\({}^2\)= \(\sum_{i=0}^{N-1} (\)x[i] - mean\() ^ 2\).

def std1D_vectorized(x, mean):
    tmp_1 = x - mean
    tmp_2 = tmp_1 ** 2
    tmp_3 = tmp_2.sum()
    return math.sqrt(tmp_3 / x.size)

И наконец, вычислив такую сумму, функция возвращает квадратный корень из нее деленной на размер массива.

Конечно, вводить переменные tmp_1, tmp_2 и tmp_3 вовсе не обязательно. Это сделано в этом примере, лишь для того, чтобы подробнее по этапам объяснить, что происходит под капотом функции ниже.

def std1D_vectorized(x, mean):
    return np.sqrt(((x - mean) ** 2).sum() / x.size)

print(x.std(), std1D_vectorized(x, x.mean()))

0.2686383564940283 0.2686383564940283

Кроме того первая версия этой функции позволяет явно обратить ваше внимание на то, что в процессе работы векторизованной функции (и изначальной и сокращенной) создаются дополнительные массивы. Для этого выделяется дополнительная память и в ряде случаев это может быть критично. Эти дополнительные массивы собираются сборщиком мусора, как только они перестают быть нужными. В первом случае на выходе из функции (когда локальные имена tmp_1, tmp_2 и tmp_3 уничтожаются) и во втором случае сразу, как только вычисления с ними закончены.