[Pyplot]使用curve_fit函数根据数据点拟合曲面

一、背景

使用python+matplotlib实现根据数据点拟合3D曲面。实现效果如图1所示：

二、代码

#!/usr/bin/env python3
import numpy as np
from scipy.optimize import curve_fit
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt


def function(data, a, b, c):
    '''
    拟合曲面使用的函数
    '''
    x = data[0]
    y = data[1]
    return a*x*x+b/y + c


# 初始化数据
x_data = np.linspace(1, 10, 20)
y_data = np.linspace(1, 10, 20)
# X, Y, Z 都是二维的数据
X_data, Y_data = np.meshgrid(x_data, y_data)

# 目标曲面的方程为 z = a*x*x + b/y + c
# 待求参数值为 [a,b,c]
Z_data = 2.0*X_data*X_data + 3.0/Y_data + 5.0

# 将数据点展开成一维
x_data = X_data.flatten()
y_data = Y_data.flatten()
z_data = Z_data.flatten()

parameters, covariance = curve_fit(function, [x_data, y_data], z_data)
# 显示根据拟合得到的曲面方程参数 [a,b,c]
print(parameters)

# 绘制结果
fig = plt.figure()
# ax = fig.add_subplot(
#     111, projection='3d')
ax = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax)
# plot surface
Z_fit = function([X_data, Y_data], *parameters)
# 绘制拟合得到的曲面
ax.plot_surface(X_data, Y_data, Z_fit, alpha=0.5)
# 绘制原始数据点
ax.scatter(x_data, y_data, z_data, color='red')
# 设置坐标轴
ax.set_xlabel('X data')
ax.set_ylabel('Y data')
ax.set_zlabel('Z data')

plt.show()