Python 3D polynomial surface fit, order dependent

Question

I am currently working with astronomical data among which I have comet images. I would like to remove the background sky gradient in these images due to the time of capture (twilight). The first program I developed to do so took user selected points from Matplotlib's "ginput" (x,y) pulled the data for each coordinate (z) and then gridded the data in a new array with SciPy's "griddata."

Since the background is assumed to vary only slightly, I would like to fit a 3d low order polynomial to this set of (x,y,z) points. However, the "griddata" does not allow for an input order:

griddata(points,values, (dimension_x,dimension_y), method='nearest/linear/cubic')

Any ideas on another function that may be used or a method for developing a leas-squares fit that will allow me to control the order?

Answer 1

Griddata uses a spline fitting. A 3rd order spline is not the same thing as a 3rd order polynomial (instead, it's a different 3rd order polynomial at every point).

If you just want to fit a 2D, 3rd order polynomial to your data, then do something like the following to estimate the 16 coefficients using all of your data points.

import itertools
import numpy as np
import matplotlib.pyplot as plt

def main():
    # Generate Data...
    numdata = 100
    x = np.random.random(numdata)
    y = np.random.random(numdata)
    z = x**2 + y**2 + 3*x**3 + y + np.random.random(numdata)

    # Fit a 3rd order, 2d polynomial
    m = polyfit2d(x,y,z)

    # Evaluate it on a grid...
    nx, ny = 20, 20
    xx, yy = np.meshgrid(np.linspace(x.min(), x.max(), nx), 
                         np.linspace(y.min(), y.max(), ny))
    zz = polyval2d(xx, yy, m)

    # Plot
    plt.imshow(zz, extent=(x.min(), y.max(), x.max(), y.min()))
    plt.scatter(x, y, c=z)
    plt.show()

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

main()