gaussfit.py
34 lines
| 1.2 KiB
| text/x-python
|
PythonLexer
Fiorella Quino
|
r178 | from numpy import * | ||
| from scipy import optimize | ||||
| def gaussian(height, center_x, center_y, width_x, width_y): | ||||
| """Returns a gaussian function with the given parameters""" | ||||
| width_x = float(width_x) | ||||
| width_y = float(width_y) | ||||
| return lambda x,y: height*exp( | ||||
| -(((center_x-x)/width_x)**2+((center_y-y)/width_y)**2)/2) | ||||
| def moments(data): | ||||
| """Returns (height, x, y, width_x, width_y) | ||||
| the gaussian parameters of a 2D distribution by calculating its | ||||
| moments """ | ||||
| total = data.sum() | ||||
| X, Y = indices(data.shape) | ||||
| x = (X*data).sum()/total | ||||
| y = (Y*data).sum()/total | ||||
| col = data[:, int(y)] | ||||
| width_x = sqrt(abs((arange(col.size)-y)**2*col).sum()/col.sum()) | ||||
| row = data[int(x), :] | ||||
| width_y = sqrt(abs((arange(row.size)-x)**2*row).sum()/row.sum()) | ||||
| height = data.max() | ||||
| return height, x, y, width_x, width_y | ||||
| def fitgaussian(data): | ||||
| """Returns (height, x, y, width_x, width_y) | ||||
| the gaussian parameters of a 2D distribution found by a fit""" | ||||
| params = moments(data) | ||||
| errorfunction = lambda p: ravel(gaussian(*p)(*indices(data.shape)) - | ||||
| data) | ||||
| p, success = optimize.leastsq(errorfunction, params) | ||||
| return p | ||||