The equivalent of `polyfit` for a two-dimensional polynomial in Python

Question

I would like to find a least squares solution for the coefficients a in

 z = (a0 + a1*x + a2*y + a3*x**2 + a4*x**2*y + a5*x**2*y**2 + a6*y**2 + a7*x*y**2 + a8*x*y) 

the given arrays are x , y and z length 20. Essentially, I'm looking for the equivalent of numpy.polyfit , but for a 2D polynomial.

This question is similar, but the solution is provided through MATLAB.

Answer 1

Here is an example showing how you can use numpy.linalg.lstsq for this task:

 import numpy as np x = np.linspace(0, 1, 20) y = np.linspace(0, 1, 20) X, Y = np.meshgrid(x, y, copy=False) Z = X**2 + Y**2 + np.random.rand(*X.shape)*0.01 X = X.flatten() Y = Y.flatten() A = np.array([X*0+1, X, Y, X**2, X**2*Y, X**2*Y**2, Y**2, X*Y**2, X*Y]).T B = Z.flatten() coeff, r, rank, s = np.linalg.lstsq(A, B) 

correction coefficients coeff :

 array([ 0.00423365, 0.00224748, 0.00193344, 0.9982576 , -0.00594063, 0.00834339, 0.99803901, -0.00536561, 0.00286598]) 

Note that coeff[3] and coeff[6] respectively correspond to X**2 and Y**2 , and they are close to 1. because sample data was created using Z = X**2 + Y**2 + small_random_component .

Answer 2

Great answer. Just add code to restore the function using the least squares solution for the coefficients a,

 def poly2Dreco(X, Y, c): return (c[0] + X*c[1] + Y*c[2] + X**2*c[3] + X**2*Y*c[4] + X**2*Y**2*c[5] + Y**2*c[6] + X*Y**2*c[7] + X*Y*c[8]) 

Answer 3

Based on the answers of @Saullo and @Francisco, I made a function that I found useful:

 def polyfit2d(x, y, z, kx=3, ky=3, order=None): ''' Two dimensional polynomial fitting by least squares. Fits the functional form f(x,y) = z. Notes ----- Resultant fit can be plotted with: np.polynomial.polynomial.polygrid2d(x, y, soln.reshape((kx+1, ky+1))) Parameters ---------- x, y: array-like, 1d x and y coordinates. z: np.ndarray, 2d Surface to fit. kx, ky: int, default is 3 Polynomial order in x and y, respectively. order: int or None, default is None If None, all coefficients up to maxiumum kx, ky, ie. up to and including x^kx*y^ky, are considered. If int, coefficients up to a maximum of kx+ky <= order are considered. Returns ------- Return paramters from np.linalg.lstsq. soln: np.ndarray Array of polynomial coefficients. residuals: np.ndarray rank: int s: np.ndarray ''' # grid coords x, y = np.meshgrid(x, y) # coefficient array, up to x^kx, y^ky coeffs = np.ones((kx+1, ky+1)) # solve array a = np.zeros((coeffs.size, x.size)) # for each coefficient produce array x^i, y^j for index, (j, i) in enumerate(np.ndindex(coeffs.shape)): # do not include powers greater than order if order is not None and i + j > order: arr = np.zeros_like(x) else: arr = coeffs[i, j] * x**i * y**j a[index] = arr.flatten() # do leastsq fitting and return leastsq result return np.linalg.lstsq(aT, np.ravel(z), rcond=None) 

And the resulting fit can be visualized with:

 fitted_surf = np.polynomial.polynomial.polygrid2d(x, y, soln.reshape((kx+1,ky+1))) plt.matshow(fitted_surf)