CODEWITHSUNDEEP
pythonnumpyscipy
Bart M
Bart M

Reputation: 737

Fitting 3D data points to polynomial surface and getting the surface equation back

I am new to Python 3D fitting, and the related optimisation techniques. I tried to understand similar topics and to find an answer based on least squares methods, but my success has been rather limited.

Problem: I have a number (ca. 400) of 3D points stored in a np array.

data = np.array([[x1, y1, z1], [x2, y2, z2], [x3, y3, z3], [x4, y4, z4], ... ])

I would like to fit a polynomial surface (order 2 or 3) to these points, and then get the parameters of the surface equation, so I could calculate z for any given pair of (x, y) values. For example:

z = (A * x ** 2) + (B * y ** 2) + (C * x * y) + (D * x) + (E * y) + F

I need to get an output that would contain the A, B, C, D, E, and F parameters. Is there a nice 'Pythonic' way of doing this?

Upvotes: 1

Views: 3200

Answers (2)

FBruzzesi
FBruzzesi

Reputation: 6505

If you allow sklearn dependency you can use sklearn.preprocessing.PolynomialFeatures together with sklearn.linear_model.LinearRegression:

import numpy as np
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression

np.random.seed(0)

# Generate fake data 
n,m = 400, 3
data = np.random.randn(n,m)

# Generate polynomial features of desired degree
d = 3
poly = PolynomialFeatures(degree=d, include_bias=False)
X = poly.fit_transform(data[:, :-1]) # X.shape = (400, 9)
y = data[:,-1] # y.shape = (400,)

# Define and fit linear regression 
clf = LinearRegression()
clf.fit(X, y)

# Check results
print(clf.coef_)
[-0.01437971  0.09894586  0.01936384  0.17245758  0.05518938
  0.0239589  -0.09930492 -0.04593238 -0.01588326]

print(clf.intercept_)
-0.12456419670821159

Upvotes: 2

Péter Leéh
Péter Leéh

Reputation: 2119

curve_fit accepts multi-dimensional array for independent variable, but your function must accept the same:

import numpy as np
from scipy.optimize import curve_fit

data = np.array(
    [[0, 0, 1],
     [1, 1, 2],
     [2, 1, 3],
     [3, 0, 5],
     [4, 0, 2],
     [5, 1, 3],
     [6, 0, 7]]
)

def func(X, A, B, C, D, E, F):
    # unpacking the multi-dim. array column-wise, that's why the transpose
    x, y, z = X.T

    return (A * x ** 2) + (B * y ** 2) + (C * x * y) + (D * x) + (E * y) + F

popt, _ = curve_fit(func, data, data[:,2])

from string import ascii_uppercase
for i, j in zip(popt, ascii_uppercase):
    print(f"{j} = {i:.3f}")

# A = 0.060
# B = 2004.446
# C = -0.700
# D = 0.521
# E = -2003.046
# F = 1.148

Note that in general you should provide initial guess for parameters to achieve good fitting results.

Upvotes: 2

Related Questions