如何在同一绘图中为多数据集绘制最佳拟合线？

一、纯NumPy + Matplotlib 解决方案 #

方法1：循环拟合每个数据集 #

这是最直观的实现方式，对每组Y数据单独进行多项式拟合，再同步绘制散点与拟合线：

import numpy as np
import matplotlib.pyplot as plt

# 示例数据（替换为你的真实数据）
x = np.arange(1, 147)  # 时间步，形状(146,)
Y = np.random.rand(4, 146)  # 4组距离数据，形状(4,146)

# 配置可视化参数
colors = ['#ff4444', '#2288ff', '#44aa44', '#ffaa00']
dataset_labels = ['数据集1', '数据集2', '数据集3', '数据集4']

plt.figure(figsize=(10,6))

for idx in range(Y.shape[0]):
    # 绘制散点图
    plt.scatter(x, Y[idx], s=12, color=colors[idx], label=dataset_labels[idx])
    # 拟合3次多项式
    poly_coeffs = np.polyfit(x, Y[idx], 3)
    # 生成密集X点让拟合线更平滑
    x_fit = np.linspace(x.min(), x.max(), 200)
    # 计算拟合Y值
    y_fit = np.polyval(poly_coeffs, x_fit)
    # 绘制拟合线
    plt.plot(x_fit, y_fit, color=colors[idx], linestyle='--', linewidth=2)

plt.xlabel('时间步')
plt.ylabel('距离')
plt.legend()
plt.title('多数据集散点图与3次多项式拟合线')
plt.show()

方法2：一次性拟合所有数据集 #

NumPy支持对多维Y数组直接拟合，只要X的长度与Y的最后一维长度匹配。调用np.polyfit后会返回形状为(4,4)的系数矩阵（每组数据对应一行，每行是3次多项式的4个系数）：

# 一次性完成4组数据的拟合
all_coeffs = np.polyfit(x, Y, 3)  # 输出形状：(4,4)

plt.figure(figsize=(10,6))
x_fit = np.linspace(x.min(), x.max(), 200)

for idx in range(Y.shape[0]):
    plt.scatter(x, Y[idx], s=12, color=colors[idx], label=dataset_labels[idx])
    # 用对应组的系数计算拟合值
    y_fit = np.polyval(all_coeffs[idx], x_fit)
    plt.plot(x_fit, y_fit, color=colors[idx], linestyle='--', linewidth=2)

plt.xlabel('时间步')
plt.ylabel('距离')
plt.legend()
plt.title('多数据集散点图与3次多项式拟合线（一次性拟合）')
plt.show()

如果之前调用np.polyfit报错，检查numpy版本（建议1.15+），旧版本对多维Y的支持有限。

二、Seaborn 简化方案 #

Seaborn更适合处理分组数据，只需先将数据转换为长格式（tidy data），即可快速完成分组拟合与绘图：

import seaborn as sns
import pandas as pd

# 将数据转换为DataFrame长格式
df = pd.DataFrame({
    '时间步': np.tile(x, 4),
    '距离': Y.flatten(),
    '数据集': np.repeat(dataset_labels, 146)
})

# 用lmplot一键实现分组拟合与绘图
g = sns.lmplot(
    data=df,
    x='时间步',
    y='距离',
    hue='数据集',
    order=3,  # 指定3次多项式
    scatter_kws={'s':12},
    line_kws={'linestyle':'--'},
    height=6,
    aspect=1.5
)
g.set_axis_labels('时间步', '距离')
g.fig.suptitle('多数据集3次多项式拟合', y=1.02)
plt.show()

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