如何将三个相似的散点图绘制函数重构为带参数的单一函数

核心思路 #

提取三个函数中重复的循环、绘图逻辑作为公共代码，把变量计算规则和图表标注信息作为参数传入通用函数，实现一次定义、多场景复用。

重构后的代码 #

通用绘图函数 #

import matplotlib.pyplot as plt
import numpy as np

def plot_correlation(x_calculator, y_calculator, x_label, y_label, title):
    x_list = []
    y_list = []
    # 循环处理ID 1-18，按传入规则计算变量值
    for l in range(1, 19):
        x_val = x_calculator(l)
        y_val = y_calculator(l)
        x_list.append(x_val)
        y_list.append(y_val)
    
    # 绘制散点图与拟合直线
    plt.scatter(x_list, y_list, color='purple')
    a, b = np.polyfit(x_list, y_list, 1)
    plt.plot(x_list, a * np.array(x_list) + b)
    
    # 设置图表标注
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.title(title)
    plt.show()

对应原三个函数的调用示例 #

1. 替代plot_RestRate_BMPmax()

plot_correlation(
    x_calculator=lambda l: test.query(f'time <0 & ID == {l} ')['BPM'].mean(),
    y_calculator=lambda l: test[test.ID == l]['BPM'].nlargest(n=5).mean(),
    x_label='Rest Rate',
    y_label='Maximum heart rate',
    title='Correlation between Maximum heart rate and Rest Rate'
)

2. 替代plot_BMP_VO2()

plot_correlation(
    x_calculator=lambda l: test[test.ID == l]['BPM'].nlargest(n=5).mean(),
    y_calculator=lambda l: test[test.ID == l]['VO2'].nlargest(n=5).mean(),
    x_label='Maximum Heart Rate',
    y_label='Maximum Oxygen Consumption',
    title='Correlation between Maximum Heart Rate and Maximum Oxygen Consumption'
)

3. 替代plot_Age_Lactate()

plot_correlation(
    x_calculator=lambda l: test[test.ID == l]['age'].values[0],
    y_calculator=lambda l: test[test.ID == l]['VO2'].nlargest(n=5).mean() * 0.8,
    x_label='Age',
    y_label='Lactate threshold',
    title='Correlation between Age and Lactate threshold'
)

可选优化说明 #

如果变量计算逻辑复杂，可以把lambda替换为单独定义的函数，提升代码可读性。比如：

def calculate_rest_rate(l):
    return test.query(f'time <0 & ID == {l} ')['BPM'].mean()

# 调用时直接传入函数名
plot_correlation(
    x_calculator=calculate_rest_rate,
    # 其他参数...
)

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