Search code examples
pythoncurve-fittinglmfit

2D Gaussian fit using lmfit


I need to fit a two dimensional Gaussian to a data set I read in. My choice of fitting routine is lmfit, as it allows easy implementation of boundary conditions and fixing of parameters. As I am not the most efficient progammer, I have problems implementing my needs. Here is what I did:

from numpy import *
from math import *
from lmfit import Parameters,minimize,report_fit

## fails to run 
# from https://www.w3resource.com/python-exercises/numpy/python-numpy-exercise-79.php
x,y = meshgrid(linspace(-1,1,10),linspace(-1,1,10))
#d = sqrt(x*x+y*y)
#sigma, mu = 1.0, 0.0
#g = exp(-( (d-mu)**2 / ( 2.0 * sigma**2 ) ) )

def gaussian2D(p,x,y):
    height = p["height"].value
    centroid_x = p["centroid_x"].value
    centroid_y = p["centroid_y"].value
    sigma_x = p["sigma_x"].value
    sigma_y = p["sigma_y"].value
    background = p["background"].value
    return height*exp(-(((centroid_x-x)/sigma_x)**2+((centroid_y-y)/sigma_y)**2)/2.0)+background

def residuals(p,x,y,z):
    return z - gaussian2D(p,x,y)

initial = Parameters()
initial.add("height",value=1.)
initial.add("centroid_x",value=0.)
initial.add("centroid_y",value=0.)
initial.add("sigma_x",value=1.)
initial.add("sigma_y",value=3.)
initial.add("background",value=0.)

xx,yy = meshgrid(x,y)

fit = minimize(residuals,initial,args=(array(xx).flatten(),array(yy).flatten(),array(g).flatten()))
popt = fit.params
print report_fit(fit)

First of all, the sample code to generate a 2D Gaussian fails to run and gives a "TypeError: only size-1 arrays can be converted to Python scalars" for d = sqrt(xx+yy). As I am using data from a file anyway, I am working with the sample data given on the website here.

Some research told me to convert the 2D arrays into 1D data in order for lmfit to be able to process them. My attempts to implement that using the flatten method on my arrays is unsuccessful, giving the same error (TypeError: only size-1 arrays can be converted to Python scalars). I am not versed enough to fully understand the code in the link.

I would appreciate any help, esp. as I prefer to define my own functions to be fitted to the data instead of relying on in-built models.


Solution

  • I think you're close, and just mixing up when (or how often) to call meshgrid. A modified version would be

    import numpy as np
    from lmfit import Parameters, minimize, report_fit
    
    x, y = np.meshgrid(np.linspace(-1, 1, 10), np.linspace(-1, 1, 10))
    
    def gaussian2D(x, y, cen_x, cen_y, sig_x, sig_y, offset):
        return np.exp(-(((cen_x-x)/sig_x)**2 + ((cen_y-y)/sig_y)**2)/2.0) + offset
    
    def residuals(p, x, y, z):
        height = p["height"].value
        cen_x = p["centroid_x"].value
        cen_y = p["centroid_y"].value
        sigma_x = p["sigma_x"].value
        sigma_y = p["sigma_y"].value
        offset = p["background"].value
        return (z - height*gaussian2D(x,y, cen_x, cen_y, sigma_x, sigma_y, offset))
    
    # test data
    g = gaussian2D(x, y, 1.2, 2.1, 0.5, 0.7, 1.1)
    
    
    initial = Parameters()
    initial.add("height",value=1.)
    initial.add("centroid_x",value=0.)
    initial.add("centroid_y",value=0.)
    initial.add("sigma_x",value=1.)
    initial.add("sigma_y",value=3.)
    initial.add("background",value=0.)
    
    
    fit = minimize(residuals, initial, args=(x, y, g))
    print(report_fit(fit))
    

    That is, define a gaussian2D() function that you can better use and test, and then have a simple objective function that just calls that.