计算两列的相关性

计算两列的相关性

使用Pandas中的corr()函数计算DataFrame中特定的两列之间相关系数。

def corr_analys(input_file_path, col_1, col_2, output_pic_path, sheet_name='Sheet1'):
    '''
    # ######################################
    # 计算两列之间的相关系数（Pearson相关系数）
    # 打印输出计算结果
    # 并绘制相关性热图
    # 输入数据建议先归一化处理
    # ######################################

    Parameters
    ----------
    input_file_path : 包含要计算的列的文件，这里是excel文件

    col_1, col_2 : 要计算的两列

    output_pic_path ：生成的相关性热图存储的路径，图片名称是日期

    sheet_name : 指定具体的excel工作表
    '''
    
    df = pd.read_excel(input_file_path, sheet_name=sheet_name)

    subset = df[[col_1, col_2]]

    corr_coefficient = subset[col_1].corr(subset[col_2])
    print('Correlation coefficient between {} and {} : {:.7f}'.format(col_1, col_2, corr_coefficient))

    corr_matrix = subset.corr()

    plt.figure(figsize=(8, 6))
    # 显示中文，防止乱码
    plt.rcParams['font.family'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    # 使用sns中的heatmap()函数讲计算出来的相关系数转化为热力图
    sns.heatmap(corr_matrix, cmap='coolwarm', annot=True, fmt='.2f')
    plt.title('Correlation Heatmap between {} and {}'.format(col_1, col_2))
    plt.xlabel(col_2)
    plt.ylabel(col_1)

    current_time = datetime.datetime.now()
    file_time = current_time.strftime("%Y%m%d%H%M%S") + '.png'
    output_file_path = output_pic_path + 'Corr_' + file_time
    # 检查目录是否存在，如果不存在则创建
    output_dir = os.path.dirname(output_pic_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    plt.savefig(output_file_path)

    plt.close()

    print('Correlation coefficient completed')

    return output_file_path


def call_corr():
    input_file_path = '../n.xlsx'
    output_pic_path = './cor/'
    col_1 = 'xxxx'
    col_2 = 'xxxx'

    output_file_path = corr_analys(input_file_path, col_1, col_2, output_pic_path)
    print('Hotmap saved to {}'.format(output_file_path))

关键的代码其实就是

corr_coefficient = subset[col_1].corr(subset[col_2])

这里corr()默认是使用Pearson方法，如果需要使用Spearman方法或者Kendall方法，需要在执行计算式显式指定。

例如下面是一个计算Spearman相关系数的函数接口：

def corr_analys_spearman(input_file_path, col_1, col_2, output_pic_path, sheet_name='Sheet1'):
    # 计算斯皮尔曼相关系数
    df = pd.read_excel(input_file_path, sheet_name=sheet_name)

    subset = df[[col_1, col_2]]

    corr_coefficient = subset[col_1].corr(subset[col_2], method='spearman')
    print('Spearman correlation coefficient between {} and {} : {:.7f}'.format(col_1, col_2, corr_coefficient))

关键是在corr中显示的指定method='spearman'

Reference

不错的网站：极客笔记

https://deepinout.com/pandas/pandas-tutorials/t_how-to-get-the-correlation-between-two-columns-in-pandas.html