如何为二维数组散点图添加带方程的线性趋势线？

我最近生成了一个8×8的二维数组，绘制了它的散点图，但在尝试添加线性趋势线并显示对应的方程时遇到了问题，也不知道怎么获取这个方程。我的代码目前是这样的（中间有重复导入和截断的错误）：

# Import Libraries
import numpy as np
import matplotlib.pyplot as plt

# Generate Random data
a = np.zeros(shape=(8,8))
a[0] = [1,2,3,4,5,6,7,8]
a[1] = [1,0,3,4,0,6,6,8]
a[2] = [1,2,3,4,5,3,7,8]
# Import Libraries（这里重复导入了，还有拼写错误impor...）
import numpy as np
impor...

别担心，我来帮你解决这个问题！要给散点图加线性趋势线并显示方程，我们需要先把二维数组的x、y数据整理成一维，再用numpy做线性拟合，最后把结果画出来并标注方程。具体步骤如下：

1. 修正代码结构，整理数据 #

首先要避免重复导入库，然后把二维数组中的x和y数据扁平化——因为散点图需要的是一维的x坐标和对应的y坐标数组。我们可以用索引生成和数组扁平化来处理：

import numpy as np
import matplotlib.pyplot as plt

# 生成你的二维数组
a = np.zeros(shape=(8,8))
a[0] = [1,2,3,4,5,6,7,8]
a[1] = [1,0,3,4,0,6,6,8]
a[2] = [1,2,3,4,5,3,7,8]

# 获取每个元素的x（列索引）和y（数组值），扁平化处理
x = np.tile(np.arange(8), 8)  # 每个行的列索引重复8次，生成所有x坐标
y = a.flatten()  # 把二维数组转成一维，得到所有y值

2. 计算线性趋势线的系数 #

用np.polyfit计算线性拟合的斜率和截距，线性拟合对应1次多项式，所以第二个参数传1：

# 线性拟合，得到系数：[斜率, 截距]
coefficients = np.polyfit(x, y, 1)
slope = coefficients[0]
intercept = coefficients[1]

# 生成趋势线的多项式函数，方便计算对应x的y值
trend_line = np.poly1d(coefficients)

3. 绘制散点图和趋势线，并显示方程 #

最后把散点图、趋势线画出来，再用plt.text把方程标注在图上：

# 绘制散点图
plt.scatter(x, y, label='Data Points')

# 绘制趋势线
plt.plot(x, trend_line(x), color='red', label='Trend Line')

# 格式化方程文本，保留两位小数让结果更清晰
equation_text = f'y = {slope:.2f}x + {intercept:.2f}'

# 把方程放在图的右上角（位置可以通过调整前两个参数自定义）
plt.text(0.8, 0.9, equation_text, transform=plt.gca().transAxes, fontsize=10)

# 添加图例和坐标轴标签
plt.xlabel('X')
plt.ylabel('Y')
plt.legend()
plt.show()

# 如果你需要直接输出方程的精确系数，也可以打印出来
print(f"线性趋势线方程：y = {slope:.4f}x + {intercept:.4f}")

如果你的二维数组里有很多初始的0值（比如后面几行都是0），这些点可能会干扰拟合结果，你可以过滤掉不需要的数据点：

# 过滤掉y=0的点
mask = y != 0
x_filtered = x[mask]
y_filtered = y[mask]

# 用过滤后的数据重新拟合
coefficients_filtered = np.polyfit(x_filtered, y_filtered, 1)
slope_filtered = coefficients_filtered[0]
intercept_filtered = coefficients_filtered[1]

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