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

Question

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.

Seymore

no sorry, I added it in there. – Seymore

sklearn has kernel ridge regression but I'm not sure whether you can set alpha to 0. Would that work for you? – Aboveground

Scipy's Splines might be useful: docs.scipy.org/doc/scipy/reference/tutorial/interpolate.html – Benzoic

The kernel ridge regression provides a nice model, but I don't see how to display the coefficients. – Seymore

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)

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.

Misericord

What is your goal with each of these lines? Z = X**2 + Y**2 + np.random.rand(*X.shape)*0.01 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 Why do you flatten this matrices? X = X.flatten() Y = Y.flatten() – Seger

@Seger the goal with the expression for Z is only to create a sample of data to test the surface fit. I had to flatten X and Y such that A becomes a 2D array, which is the format needed by the lstsq solvers. – Misericord

Answer 2

Based on the answers from @Saullo and @Francisco I have made a function which I have found helpful:

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.ravel()

    # do leastsq fitting and return leastsq result
    return np.linalg.lstsq(a.T, np.ravel(z), rcond=None)

And the resultant fit can be visualised with:

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

Teresitateressa

The suggested edit queue for this answer is full which sounds like there are over 500 edits people have tried to submit. For anyone else who comes across this, polyval2d in the bottom portion should be polygrid2d as noted in the comments. soln is the first portion of the return value, again as noted in the comments. I'd also suggest putting full code to call your code/ plot, but that's too much for a comment. I really appreciate your answer, I just think it might be easier for people to use if there was a more robust usage example. – Duero

I also believe the line for index, (j, i) in enumerate(np.ndindex(coeffs.shape)): should be for index, (i, j) in enumerate(np.ndindex(coeffs.shape)): since values where kx!=ky throws an error – Duero

Answer 3

Excellent answer by Saullo Castro. Just to add the code to reconstruct the function using the least-squares solution for the a coefficients,

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])

Polymorphous

Answer 4

You can also use scikit-learn for this.

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

x = np.linspace(0, 1, 20)
y = np.linspace(0, 1, 20)
X, Y = np.meshgrid(x, y, copy=False)
X = X.flatten()
Y = Y.flatten()

# Generate noisy data
np.random.seed(0)
Z = X**2 + Y**2 + np.random.randn(*X.shape)*0.01

# Process 2D inputs
poly = PolynomialFeatures(degree=2)
input_pts = np.stack([X, Y]).T
assert(input_pts.shape == (400, 2))
in_features = poly.fit_transform(input_pts)

# Linear regression
model = LinearRegression(fit_intercept=False)
model.fit(in_features, Z)

# Display coefficients
print(dict(zip(poly.get_feature_names_out(), model.coef_.round(4))))

# Check fit
print(f"R-squared: {model.score(poly.transform(input_pts), Z):.3f}")

# Make predictions
Z_predicted = model.predict(poly.transform(input_pts))

Out:

{'1': 0.0012, 'x0': 0.003, 'x1': -0.0074, 'x0^2': 0.9974, 'x0 x1': 0.0047, 'x1^2': 1.0014}
R-squared: 1.000

Notes

I used the fit_intercept=False argument when defining the linear regression model because the polynomial features by default include the bias term '1'. Alternatively, you could disable that using

poly = PolynomialFeatures(degree=2, include_bias=False)

and then use a regular LinearRegression model with an intercept coefficient:

model.intercept_  # 0.0011741137800872492

Shoffner

get_feature_names_out() -> get_feature_names() github.com/scikit-learn-contrib/category_encoders/issues/348 - didn't follow all the discussion but dropping the _out makes it work in my version – Ked

The code still works for me in scikit-learn version 1.2.1. What version are you using? – Shoffner

My sklearn.__version__ is 0.23.2, 0.23.2-5ubuntu6 apt installed in ubuntu 22.04 - so the change is going the other way, older versions don't have the _out() – Ked

Answer 5

Note that if kx != ky the code will fail because the j and i indices are inverted in the loop.

You get (j,i) from enumerate(np.ndindex(coeffs.shape)), but then you address elements in coeffs as coeffs[i,j]. Since the shape of the coefficient matrix is given by the maximum polynomial order that you are asking to use, the matrix will be rectangular if kx != ky and you will exceed one of its dimensions.

Lip