import numpy as np
from scipy import optimize# 最小二乘法拟合
import matplotlib.pyplot as plt # python matplotlib 绘图
from mpl_toolkits.mplot3d import Axes3D # 3D 绘图
def func(x, y, p):
""" 数据拟合所用的函数：z=ax+by:param x: 自变量 x:param y: 自变量 y:param p: 拟合参数 a, b"""
a, b = p
return a * x + b * y
def residuals(p, z, x, y):
""" 得到数据 z 和拟合函数之间的差"""
return z - func(x, y, p)
def main():
x = np.arange(5)
y = np.arange(5)
z = np.array([2, 4, 7, 7, 10]) # 数据随便取的
plsq = optimize.leastsq(residuals, np.array([0, 0]), args=(z, x, y)) # 最小二乘法拟合
# [0, 0] 为参数 a, b 初始值
a, b = plsq[0] # 获得拟合结果
print("拟合结果:\na ={}".format(a))
print("b ={}".format(b))
# 绘图
xp = np.linspace(-1, 6, 100)
yp = np.linspace(-1, 6, 100)
X, Y = np.meshgrid(xp, yp)
Z = func(X, Y, [a, b]) # 带入拟合得到的 a, b
fig = plt.figure(figsize=(8, 6))
ax = Axes3D(fig) # 3D 绘图
ax.plot_surface(X, Y, Z, alpha=0.5)
ax.scatter(x, y, z, color="r")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
plt.show()
if __name__ == "__main__":
main()