如何在for循环中为stats.linregress设置r值阈值过滤无效变量组合

原有代码核心问题 #

• 自行维护的变量去重逻辑存在缺陷，会重复添加相同变量、漏掉有效组合，且双重循环会重复计算(a,b)和(b,a)两组完全等价的组合，额外浪费算力
• 先绘图再计算回归统计量，即使组合不符合阈值要求也已经生成了绘图对象，容易出现空白图、报错问题

修改后代码 #

import itertools
import seaborn as sns
import matplotlib.pyplot as plt
from scipy import stats

# 可自定义阈值：r绝对值最小值、p值最大值
R_ABS_MIN = 0.3
P_MAX = 0.05

# 直接生成所有不重复的两两变量组合，自动避免重复计算和自身对自身的组合
for column_a, column_b in itertools.combinations(numeric_variable_list, 2):
    # 先计算回归统计量，不符合条件直接跳过，不执行绘图逻辑
    m, b, r, p, err = stats.linregress(df[column_a], df[column_b])
    if abs(r) < R_ABS_MIN or p > P_MAX:
        continue
    
    # 只有符合阈值的组合才执行绘图逻辑
    correlation_plot = sns.regplot(x=df[column_a], y=df[column_b])
    sns.despine(top=True, right=True)
    # 拼接标注文本
    textstr= 'y  = ' + str(round(m, 2)) + 'x + ' + str(round(b, 2)) + '\n'
    textstr += 'r2 = ' + str(round(r**2, 4)) + '\n'
    textstr += 'p  = ' + str(round(p, 4))
    plt.text(0.15, 0.70, textstr, fontsize=12, transform=plt.gcf().transFigure)
    plt.show(correlation_plot)
    plt.close()

优化说明 #

• 用itertools.combinations生成无重复的两两变量组合，100个变量的场景下循环次数直接从10000次降到4950次，算力消耗减半
• 先计算回归统计量再做阈值判断，不符合要求的组合直接跳过绘图步骤，完全避免空白图和多余的资源消耗
• 阈值可根据业务需求自行调整，比如只需要过滤p值就保留p值判断条件，只需要过滤r值就保留r值判断条件即可

内容的提问来源于stack exchange，提问作者dvwut