Корреляция Спирмена со столбцами матрицы: spearmanr, NumPy и обработка связок

Почему spearmanr даёт квадратную матрицу и как получить постолбцовую корреляцию Спирмена быстро: срез нужного столбца, ранги в NumPy, точный учёт связок.

Когда нужно посчитать корреляцию Спирмена между опорным вектором и каждым столбцом матрицы, очевидный цикл работает, но медленно. Попытка «векторизовать» расчёт, подав в scipy.stats.spearmanr целые массивы, часто неожиданно возвращает большую квадратную матрицу. Разберёмся, почему так происходит и как получить именно постолбцовые корреляции эффективно — включая вариант на чистом NumPy, который избегает лишних попарных вычислений.

Постановка задачи

Пусть у вас есть матрица данных с 37 наблюдениями и N столбцами (образцами), а также одномерный опорный вектор из тех же 37 наблюдений. Цель — вычислить корреляцию Спирмена между каждым столбцом и опорным вектором, получив результат формы (N, 1).

Ниже — простой вариант, который даёт правильный ответ:

from scipy.stats import spearmanr
import numpy as np
prng = np.random.default_rng()
X = prng.standard_normal((37, 100))
ref_vec = prng.standard_normal((37, 1))
corr_list = []
for col in X.T:
    corr_list.append(spearmanr(col, ref_vec).statistic)
result = np.array(corr_list)
print(result.shape)  # (100,)

Он делает именно то, что нужно, но обрабатывает данные по одному столбцу. Попытка вызвать spearmanr на целых массивах поначалу выглядит многообещающе:

from scipy.stats import spearmanr
import numpy as np
prng = np.random.default_rng()
X = prng.standard_normal((37, 100))
ref_vec = prng.standard_normal((37, 1))
pairwise = spearmanr(X, ref_vec)
print(pairwise.statistic.shape)  # (101, 101)

Вместо вектора (N, 1) вы получаете матрицу (N+1, N+1). Это не ошибка — так устроен API для двумерных входов.

Почему размерность «раздувается» у spearmanr на 2D-входах

В 1D-случае spearmanr(x, y) возвращает одно число — корреляцию. В 2D-случае spearmanr(A) воспринимает каждый столбец A как отдельную переменную и вычисляет полную попарную матрицу корреляций между столбцами. Передача двух двумерных массивов эквивалентна конкатенации их столбцов с последующим вычислением попарных корреляций между всеми получившимися столбцами. Поэтому при подаче X со 100 столбцами и ref_vec с 1 столбцом получается матрица 101×101: считаются все корреляции между 101 переменной.

Иными словами, вы просили всего N корреляций, а посчитали все N×N внутри X, плюс всё, что связано с опорным столбцом, и ещё симметричные элементы. На деле вам нужны только корреляции между каждым столбцом X и опорным вектором.

Векторизованный результат без лишних циклов

Поскольку опорный столбец в концептуальной конкатенации идёт последним, нужную вам колонку можно просто извлечь из большой матрицы. Пусть n — число столбцов X. Искомые постолбцовые корреляции Спирмена находятся в последнем столбце, ограниченном первыми n строками.

from scipy.stats import spearmanr
import numpy as np
prng = np.random.default_rng()
X = prng.standard_normal((37, 100))
ref_vec = prng.standard_normal((37, 1))
n = X.shape[1]
full_corr = spearmanr(X, ref_vec).statistic
colwise_corr = full_corr[:n, n]  # форма: (n,)

Так вы избегаете ручных Python-циклов, и для умеренных N это быстро, хотя при этом вычисляется гораздо больше корреляций, чем в итоге используется. В размерах выше посчитать полную матрицу 101×101 и вырезать последний столбец быстрее, чем идти столбец за столбцом. Но учтите: при очень больших N цена полной попарной матрицы становится избыточной.

Можно подумать о np.vectorize ради «красивого» интерфейса, но ускорения это не даст: по скорости это близко к обычному циклу.

from scipy.stats import spearmanr
import numpy as np
prng = np.random.default_rng()
X = prng.standard_normal((37, 100))
ref_vec = prng.standard_normal((37, 1))
def one_corr(u, v):
    return spearmanr(u, v).statistic
vec_apply = np.vectorize(one_corr, signature='(m),(m,1)->()')
colwise_corr = vec_apply(X.T, ref_vec)  # по скорости сопоставимо с циклом

Чистый NumPy: Спирмен через ранги (без совпадающих значений)

Корреляция Спирмена — это корреляция Пирсона, посчитанная по рангам. Если данные непрерывные и связки (ties) не ожидаются, ранги можно быстро получить «двойным argsort», а затем использовать закрытую форму для среднего и дисперсии рангов. Такой подход избегает большой попарной матрицы и векторизован по всем столбцам.

import numpy as np
def spearman_fast_noties(A, b):
    # A: (m, n), b: (m, 1)
    m = b.shape[0]
    mean_r = (m - 1) / 2
    var_r = (m - 1) * (m + 1) / 12
    rA = A.argsort(axis=0).argsort(axis=0)             # ранги 0..m-1
    rb = b.argsort(axis=0).argsort(axis=0)             # ранги 0..m-1
    # Пирсон по рангам с известными средним и дисперсией
    return ((rA * rb).mean(axis=0) - mean_r**2) / var_r
# Пример
prng = np.random.default_rng()
X = prng.standard_normal((37, 100))
ref_vec = prng.standard_normal((37, 1))
colwise_corr = spearman_fast_noties(X, ref_vec)

Этот вариант и векторизован, и не выполняет лишних попарных вычислений. В показанных размерах он заметно быстрее цикла: вычисление рангов через argsort дешёвое при малом m, и всё выполняется внутри NumPy.

Корректная обработка связок (повторов)

Если связки возможны и их нужно учесть точно, используйте scipy.stats.rankdata для вычисления рангов со обработкой связок. В этом случае нельзя опираться на закрытую форму среднего и дисперсии рангов: вычислите их по ранжированным массивам и примените обычную формулу Пирсона к рангам.

import numpy as np
from scipy.stats import rankdata
def spearman_tie_safe(A, b):
    # A: (m, n), b: (m, 1)
    rA = rankdata(A, axis=0)
    rb = rankdata(b, axis=0)
    # Пирсон по рангам с явным центрированием и масштабированием
    return ((rA * rb).mean(axis=0) - rA.mean(axis=0) * rb.mean()) / (rA.std(axis=0) * rb.std())
# Пример
prng = np.random.default_rng()
X = prng.standard_normal((37, 100))
ref_vec = prng.standard_normal((37, 1))
colwise_corr = spearman_tie_safe(X, ref_vec)

Этот подход векторизован и корректен при наличии связок. Для малых m rankdata может быть медленнее «двойного argsort», но гарантирует правильную обработку связок.

Почему это важно

Понимание того, как spearmanr интерпретирует 2D-входы, помогает избежать путаницы и лишних вычислений. Если вы подаёте несколько столбцов, spearmanr возвращает полную попарную матрицу корреляций по всем столбцам конкатенированных входов. Это бывает полезно, но когда нужны корреляции только с одним опорным вектором, это избыточно. Зная, где находится нужный столбец, его можно просто вырезать и тем самым избежать ручных циклов. Когда число наблюдений невелико, а число столбцов большое, работа напрямую с рангами может быть значительно быстрее и вычисляет только то, что действительно требуется.

Итоги

Если вызываете spearmanr на матрице и опорном векторе, берите нужный столбец из последнего столбца результата срезом — это эффективно. Если хотите вовсе избежать попарной матрицы, посчитайте ранги сами и примените Пирсона к рангам. Когда связки важны, используйте rankdata; когда нет — приём с двойным argsort даёт эффективную альтернативу. Избегайте np.vectorize ради производительности: он ведёт себя как обычный Python-цикл. С этими приёмами вы получаете предсказуемые формы, меньше сюрпризов и решение, которое масштабируется под вашу структуру данных.