Linear least squares fitting of a two-dimensional data


This earlier blog post presented a way of performing a non-linear least squares fit on two-dimensional data using a sum of (2D) Gaussian functions.

Fitting a two-dimensional polynomial to a surface is, in principle, a linear least-squares problem, since the fitting function is linear in the fit coefficients, $c_{i,j}$: $$ z_\mathrm{fit}(x, y) = c_{0,0} + c_{1,0}x + c_{0, 1}y + c_{2,0} x^2 + c_{1,1} xy + c_{0,2}y^2 + \ldots $$

The code below demonstrates the process, using NumPy's linalg.lstsq method.

The 2D function to be fit: a sum of two Gaussian functions with synthetic noise added:

The 2D function to fit

The fitted polynomial function and residuals plotted on a plane under the fitted data:

Fitted 2D function

The result in 2D with the fitted data contours superimposed on the noisy data:

Fitted 2D function as contours

Disclaimer: The code below is probably numerically rather unstable and should probably not be used by anyone for any purpose, especially with max_order > 3 or so.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


# The two-dimensional domain of the fit.
xmin, xmax, nx = -5, 4, 75
ymin, ymax, ny = -3, 7, 150
x, y = np.linspace(xmin, xmax, nx), np.linspace(ymin, ymax, ny)
X, Y = np.meshgrid(x, y)

# Our function to fit is going to be a sum of two-dimensional Gaussians
def gaussian(x, y, x0, y0, xalpha, yalpha, A):
    return A * np.exp( -((x-x0)/xalpha)**2 -((y-y0)/yalpha)**2)

# A list of the Gaussian parameters: x0, y0, xalpha, yalpha, A
gprms = [(0, 2, 2.5, 5.4, 1.5),
         #(-1, 4, 6, 2.5, 1.8),
         (-3, -0.5, 1, 2, 4),
         #(3, 0.5, 2, 1, 5)

# Standard deviation of normally-distributed noise to add in generating
# our test function to fit.
noise_sigma = 0.1

# The function to be fit is Z.
Z = np.zeros(X.shape)
for p in gprms:
    Z += gaussian(X, Y, *p)
Z += noise_sigma * np.random.randn(*Z.shape)

# Plot the 3D figure of the fitted function and the residuals.
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot_surface(X, Y, Z, cmap='viridis')

def get_basis(x, y, max_order=4):
    """Return the fit basis polynomials: 1, x, x^2, ..., xy, x^2y, ... etc."""
    basis = []
    for i in range(max_order+1):
        for j in range(max_order - i +1):
            basis.append(x**j * y**i)
    return basis

# We need to ravel the meshgrids of X, Y points to a pair of 1-D arrays.
x, y = X.ravel(), Y.ravel()
# Maximum order of polynomial term in the basis.
max_order = 8
basis = get_basis(x, y, max_order)
# Linear, least-squares fit.
A = np.vstack(basis).T
b = Z.ravel()
c, r, rank, s = np.linalg.lstsq(A, b, rcond=None)

print('Fitted parameters:')

# Calculate the fitted surface from the coefficients, c.
fit = np.sum(c[:, None, None] * np.array(get_basis(X, Y, max_order))
                .reshape(len(basis), *X.shape), axis=0)

rms = np.sqrt(np.mean((Z - fit)**2))
print('RMS residual =', rms)

# Plot the 3D figure of the fitted function and the residuals.
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot_surface(X, Y, fit, cmap='viridis')
cset = ax.contourf(X, Y, Z-fit, zdir='z', offset=-4, cmap='viridis')

# Plot the test data as a 2D image and the fit as overlaid contours.
fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(Z, origin='lower', cmap='viridis',
          extent=(x.min(), x.max(), y.min(), y.max()))
ax.contour(X, Y, fit, colors='w')
Current rating: 5


Comments are pre-moderated. Please be patient and your comment will appear soon.

There are currently no comments

New Comment


required (not published)