Source code for duo.core.surface_fitting

import numpy as np
import itertools

#------------------------------------------------------------
#------------------------------------------------------------
def PolyFit2D(x, y, z, order = 3):
    ncols = (order + 1) ** 2
    G = np.zeros((x.size, ncols))
    ij = itertools.product(range(order + 1), range(order + 1))
    for k, (i,j) in enumerate(ij):
        G[:,k] = x ** i * y ** j
    m, _, _, _ = np.linalg.lstsq(G, z)
    return m

#------------------------------------------------------------
#------------------------------------------------------------
def PolyVal2D(x, y, m):
    order = int(np.sqrt(len(m))) - 1
    ij = itertools.product(range(order + 1), range(order + 1))
    z = np.zeros_like(x)
    for a, (i,j) in zip(m, ij):
        z += a * x ** i * y ** j
    return z