pythonmathnumpylinear-algebrapolynomial-math

# Equivalent of `polyfit` for a 2D polynomial in Python

I'd like to find a least-squares solution for the `a` coefficients 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)
``````

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

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

Solution

• 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)
``````

the adjusting coefficients `coeff` are:

``````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 the example data was created with `Z = X**2 + Y**2 + small_random_component`.