Equivalente a `polyfit` para un polinomio 2D en Python

Me gustaría encontrar una solución de mínimos cuadrados para los coeficientes en

 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) 

dados arrays x , y , z de longitud 20. Básicamente estoy buscando el equivalente de numpy.polyfit pero para un polinomio 2D.

Esta pregunta es similar, pero la solución se proporciona a través de MATLAB.

Aquí hay un ejemplo que muestra cómo puede usar numpy.linalg.lstsq para esta tarea:

 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) 

Los coeficientes de ajuste coeff son:

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

Tenga en cuenta que coeff[3] y coeff[6] respectivamente corresponden a X**2 e Y**2 , y están cerca de 1. porque los datos de ejemplo se crearon con Z = X**2 + Y**2 + small_random_component .