三维散点图拟合仅得抛物线，如何实现平面拟合？

问题描述 #

在三维坐标系中绘制数据散点图时，当前使用二次曲面拟合函数得到了抛物线形状的结果，但实际需要平面拟合效果，修改拟合方程时出现代码运行错误，无法排查。原代码如下：

def func(xy, a, b, c, d, e, f): 
    x, y = xy 
    return a + b*x + c*y + d*x**2 + e*y**2 + f*x*y

x = temp[xStr].to_numpy()
y = temp[yStr].to_numpy()
z = temp[zStr].to_numpy()

popt, pcov = curve_fit(func, (x, y), z) 

fig = plt.figure(figsize = (16, 9))
ax = fig.add_subplot(111, projection='3d') 

ax.set_xlim(0.15, 0.55)
ax.set_ylim(0, 35)
ax.set_zlim(60, 95)


x_range = np.linspace(0, 1, 50) 
y_range = np.linspace(0, 1, 50)

X, Y = np.meshgrid(x_range, y_range) 
Z = func((X, Y), *popt) 

ax.plot_surface(X, Y, Z, color='red', alpha=1) 
 

my_cmap = plt.get_cmap('cool')
 

sctt = ax.scatter3D(x, y, z,
                    alpha = 0.8,
                    c = (z), 
                    cmap = my_cmap, 
                    marker ='^')
 
plt.title("Line " + str(arg1))
ax.set_xlabel('Head CU', fontweight ='bold') 
ax.set_ylabel('ASCU', fontweight ='bold') 
ax.set_zlabel('Recovery', fontweight ='bold')
fig.colorbar(sctt, ax = ax, shrink = 0.5, aspect = 5)

解决方案 #

要实现平面拟合，核心是将拟合函数改为线性平面方程（无二次项），同时调整相关参数匹配新函数，具体修改如下：

1. 修改拟合函数 #

平面的数学表达式为 z = a + b*x + c*y，因此将func简化为仅包含一次项和常数项：

def func(xy, a, b, c): 
    x, y = xy 
    return a + b*x + c*y

2. 修正网格范围 #

原代码中x_range和y_range设置为0-1，与数据实际范围（x:0.15-0.55，y:0-35）不匹配，导致拟合平面可能偏离散点区域，需调整为对应数据范围：

x_range = np.linspace(0.15, 0.55, 50) 
y_range = np.linspace(0, 35, 50)

3. 完整修改后代码 #

import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

# 定义平面拟合函数
def func(xy, a, b, c): 
    x, y = xy 
    return a + b*x + c*y

# 提取数据（假设temp、xStr、yStr、zStr已定义）
x = temp[xStr].to_numpy()
y = temp[yStr].to_numpy()
z = temp[zStr].to_numpy()

# 执行平面拟合
popt, pcov = curve_fit(func, (x, y), z) 

# 创建3D绘图
fig = plt.figure(figsize = (16, 9))
ax = fig.add_subplot(111, projection='3d') 

# 设置坐标轴范围
ax.set_xlim(0.15, 0.55)
ax.set_ylim(0, 35)
ax.set_zlim(60, 95)

# 生成匹配数据范围的网格
x_range = np.linspace(0.15, 0.55, 50) 
y_range = np.linspace(0, 35, 50)
X, Y = np.meshgrid(x_range, y_range) 

# 计算拟合平面的Z值
Z = func((X, Y), *popt) 

# 绘制拟合平面
ax.plot_surface(X, Y, Z, color='red', alpha=0.5) 

# 绘制散点图
my_cmap = plt.get_cmap('cool')
sctt = ax.scatter3D(x, y, z,
                    alpha = 0.8,
                    c = (z), 
                    cmap = my_cmap, 
                    marker ='^')
 
# 设置标题和坐标轴标签
plt.title("Line " + str(arg1))
ax.set_xlabel('Head CU', fontweight ='bold') 
ax.set_ylabel('ASCU', fontweight ='bold') 
ax.set_zlabel('Recovery', fontweight ='bold')
fig.colorbar(sctt, ax = ax, shrink = 0.5, aspect = 5)

plt.show()

关键改动说明 #

• 移除了原函数中的二次项（d*x**2、e*y**2、f*x*y），确保拟合结果为平面而非曲面。
• 调整网格生成范围，让拟合平面完全覆盖散点数据的分布区域，提升可视化效果。
• 降低了平面的alpha值（从1改为0.5），避免遮挡散点，更直观对比拟合效果。

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