Fitting higher order 2D surface polynomials to data

fit_poly_surface_2d.py

import numpy as np
import itertools
 
def _xy_powers(order):
    """Get powers of an x-y polynomial of certain order

    Order = 2 will have (m, n) for x^m * y^n with m+n <= order
    Example:
        >>> _xy_powers(2)
        [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (2, 0)]
    """
    powers = itertools.product(range(order + 1), range(order + 1))
    return [tup for tup in powers if sum(tup) <= order]


def polyfit2d(x, y, z, order=3):
    ncols = (order + 1)**2
    G = np.zeros((x.size, ncols))
    ij = _xy_powers(order)
    for k, (i, j) in enumerate(ij):
        G[:, k] = x**i * y**j
    m, _, _, _ = np.linalg.lstsq(G, z, rcond=None)
    return m


def polyval2d(x, y, z, m):
    order = int(np.sqrt(len(m))) - 1
    z_out = np.zeros_like(z)
    ij = _xy_powers(order)
    for a, (i, j) in zip(m, ij):
        z_out += a * x**i * y**j
    return z_out

def fit_surface(z):
    yidxs, xidxs = matrix_indices(z.shape, flatten=True)
    m = polyfit2d(xidxs, yidxs, z.flatten(), order=order)
    y_block, x_block = matrix_indices(z.shape, flatten=False)
    return polyval2d(x_block, y_block, z, m)