Skip to content

Instantly share code, notes, and snippets.

@keysie
Created August 29, 2022 07:45
Show Gist options
  • Save keysie/fc747ed6198acd3cf005567102ab5230 to your computer and use it in GitHub Desktop.
Save keysie/fc747ed6198acd3cf005567102ab5230 to your computer and use it in GitHub Desktop.
2-dimensional function fitting using numpy
import numpy as np
import matplotlib.pyplot as plt
x = np.array([6,8,10,12,15])
y = np.array([2000, 1500, 1000])
X, Y = np.meshgrid(x, y, copy=False)
Z = np.array([[1945, 2980, 4085, 5317, 7360],
[1525, 2300, 3215, 4145, 5744],
[1127, 1700, 2330, 3037, 4235]])
X = X.flatten()
Y = Y.flatten()
A = np.array([X*0+1, X, Y, X**2, X*Y]).T
B = Z.flatten()
coeff, r, rank, s = np.linalg.lstsq(A, B)
print(coeff)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
xdisp = np.linspace(6,15,100)
ydisp = np.linspace(1000,2000,100)
Xdisp, Ydisp = np.meshgrid(xdisp, ydisp, copy=False)
Adisp = np.array([Xdisp*0+1, Xdisp, Ydisp, Xdisp**2, Xdisp*Ydisp]).T
Zdisp = np.matmul(Adisp, coeff).T
ax.scatter(X,Y,Z, c='red', marker='x')
ax.plot_surface(Xdisp, Ydisp, Zdisp)
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment