Как рассчитать коэффициент корреляции сгруппированных величин в Pandas?

Question

У меня есть DataFrame в котором каждый ряд представляет дорожно-транспортное происшествие. Два из столбцов Speed_limit а также Number_of_casualties, Я хотел бы вычислить коэффициент корреляции Пирсона между ограничением скорости и отношением числа жертв к несчастным случаям для каждого ограничения скорости.

Мое решение до сих пор состоит в том, чтобы получить соответствующие количества в виде массивов и использовать SciPy's pearsonr:

import pandas as pd
import scipy.stats

df = pd.DataFrame({'Speed_limit': [10, 10, 20, 20, 20, 30],
                   'Number_of_casualties': [1, 2, 3, 4, 1, 4]})

accidents_per_speed_limit = df['Speed_limit'].value_counts().sort_index()

number_of_casualties_per_speed_limit = df.groupby('Speed_limit').sum()['Number_of_casualties']

speed_limit = accidents_per_speed_limit.index
ratio = number_of_casualties_per_speed_limit.values / accidents_per_speed_limit.values

r, _ = scipy.stats.pearsonr(x=speed_limit, y=ratio)

print("The Pearson's correlation coefficient between the number of casualties per accidents and the speed limit is {r}.".format(r=r))

Однако мне кажется, что это можно сделать более элегантно, используя метод pandas.DataFrame.corr. Как я могу реорганизовать этот код, чтобы сделать его более pandas-лайк?

Answer 1

Вместо подсчета и суммы вы можете использовать напрямую mean групповых данных затем использовать series corr(по умолчанию используется метод Pearson), т.е.

m = df.groupby('Speed_limit').mean().reset_index()
m['Speed_limit'].corr(m['Number_of_casualties'])

Выход:

0.99926008128973687

Answer 2

Я нашел следующий способ, используя два вспомогательных DataFrames:

df_aux = df.groupby('Speed_limit').agg(['count', 'sum'])
df_aux2 = pd.DataFrame({'ratio': df_aux['Number_of_casualties', 'sum'] / df_aux['Number_of_casualties', 'count'],
                        'speed_limit': df_aux.index})
print(df_aux2.corr()['ratio']['speed_limit'])

что подтверждает результат, полученный с scipy.stats.pearsonr, Это все еще не очень элегантно, хотя, и я был бы признателен за предложения по улучшению.