import numpy
import scipy.optimize
import scipy.special
import inspect
import re
[docs]class FitBase(object):
parameters = None
guess = None
result = None
summary = None
fitdata = None
def __init__(self, space, guess=None):
self.space = space
code = inspect.getsource(self.func)
args = tuple( re.findall('\((.*?)\)', line)[0].split(',') for line in code.split('\n')[2:4])
if space.dimension != len(args[0]):
raise ValueError('dimension mismatch: space has {0}, {1.__class__.__name__} expects {2}'.format(space.dimension, self, len(args[0])))
self.parameters = args[1]
self.xdata, self.ydata, self.cxdata, self.cydata = self._prepare(self.space)
if guess is not None:
if len(guess) != len(self.parameters):
raise ValueError('invalid number of guess parameters {0!r} for {1!r}'.format(guess, self.parameters))
self.guess = guess
else:
self._guess()
self.success = self._fit()
@staticmethod
def _prepare(space):
ydata = space.get_masked()
cydata = ydata.compressed()
imask = ~ydata.mask
xdata = space.get_grid()
cxdata = tuple(d[imask] for d in xdata)
return xdata, ydata, cxdata, cydata
def _guess(self):
# the implementation should use space and/or self.xdata/self.ydata and/or the cxdata/cydata maskless versions to obtain guess
raise NotImplementedError
def _fitfunc(self, params):
return self.cydata - self.func(self.cxdata, params)
def _fit(self):
result = scipy.optimize.leastsq(self._fitfunc, self.guess, full_output=True, epsfcn=0.000001)
self.message = re.sub('\s{2,}', ' ', result[3].strip())
self.result = result[0]
errdata = result[2]['fvec']
if result[1] is None:
self.variance = numpy.zeros(len(self.result))
else:
self.variance = numpy.diagonal(result[1] * (errdata**2).sum() / (len(errdata) - len(self.result)))
self.fitdata = numpy.ma.array(self.func(self.xdata, self.result), mask=self.ydata.mask)
self.summary = '\n'.join('%s: %.4g +/- %.4g' % (n, p, v) for (n, p, v) in zip(self.parameters, self.result, self.variance))
return result[4] in (1, 2, 3, 4) # corresponds to True on success, False on failure
def __str__(self):
return '{0.__class__.__name__} fit on {1}\n{2}\n{3}'.format(self, self.space, self.message, self.summary)
[docs]class PeakFitBase(FitBase):
def __init__(self, space, guess=None, loc=None):
if loc != None:
self.argmax = tuple(loc)
else:
self.argmax = None
super(PeakFitBase, self).__init__(space, guess)
def _guess(self):
maximum = self.cydata.max() # for background determination
background = self.cydata < (numpy.median(self.cydata) + maximum) / 2
if any(background == True): # the fit will fail if background is flas for all
linparams = self._linfit(list(grid[background] for grid in self.cxdata), self.cydata[background])
else:
linparams = numpy.zeros(len(self.cxdata) + 1)
simbackground = linparams[-1] + numpy.sum(numpy.vstack(param * grid.flatten() for (param, grid) in zip(linparams[:-1], self.cxdata)), axis=0)
signal = self.cydata - simbackground
if self.argmax != None:
argmax = self.argmax
else:
argmax = tuple((signal * grid).sum() / signal.sum() for grid in self.cxdata)
argmax_bkg = linparams[-1] + numpy.sum(numpy.vstack(param * grid.flatten() for (param, grid) in zip(linparams[:-1], argmax)))
try:
maximum = self.space[argmax] - argmax_bkg
except ValueError:
maximum = self.cydata.max()
if numpy.isnan(maximum):
maximum = self.cydata.max()
self.set_guess(maximum, argmax, linparams)
def _linfit(self, coordinates, intensity):
coordinates = list(coordinates)
coordinates.append(numpy.ones_like(coordinates[0]))
matrix = numpy.vstack(coords.flatten() for coords in coordinates).T
return numpy.linalg.lstsq(matrix, intensity)[0]
[docs]class AutoDimensionFit(FitBase):
def __new__(cls, space, guess=None):
if space.dimension in cls.dimensions:
return cls.dimensions[space.dimension](space, guess)
else:
raise TypeError('{0}-dimensional space not supported for {1.__name__}'.format(space.dimension, cls))
# utility functions
[docs]def rot2d(x, y, th):
xrot = x * numpy.cos(th) + y * numpy.sin(th)
yrot = - x * numpy.sin(th) + y * numpy.cos(th)
return xrot, yrot
[docs]def rot3d(x, y, z, th, ph):
xrot = numpy.cos(th) * x + numpy.sin(th) * numpy.sin(ph) * y + numpy.sin(th) * numpy.cos(ph) * z
yrot = numpy.cos(ph) * y - numpy.sin(ph) * z
zrot = -numpy.sin(th) * x + numpy.cos(th) * numpy.sin(ph) * y + numpy.cos(th) * numpy.cos(ph) * z
return xrot, yrot, zrot
[docs]def get_class_by_name(name):
options = {}
for k, v in globals().items():
if isinstance(v, type) and issubclass(v, FitBase):
options[k.lower()] = v
if name.lower() in options:
return options[name.lower()]
else:
raise ValueError("unsupported fit function '{0}'".format(name))
# fitting functions
[docs]class Lorentzian1D(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x, ) = grid
(I, loc, gamma, slope, offset) = params
return I / ((x - loc)**2 + gamma**2) + offset + x * slope
[docs] def set_guess(self, maximum, argmax, linparams):
gamma0 = 5 * self.space.axes[0].res # estimated FWHM on 10 pixels
self.guess = [maximum, argmax[0], gamma0, linparams[0], linparams[1]]
[docs]class Lorentzian1DNoBkg(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x, ) = grid
(I, loc, gamma) = xxx_todo_changeme3
return I / ((x - loc)**2 + gamma**2)
[docs] def set_guess(self, maximum, argmax, linparams):
gamma0 = 5 * self.space.axes[0].res # estimated FWHM on 10 pixels
self.guess = [maximum, argmax[0], gamma0]
[docs]class PolarLorentzian2Dnobkg(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x, y) = grid
(I, loc0, loc1, gamma0, gamma1, th) = params
a, b = tuple(grid - center for grid, center in zip(rot2d(x, y, th), rot2d(loc0, loc1, th)))
return (I / (1 + (a / gamma0)**2 + (b / gamma1)**2))
[docs] def set_guess(self, maximum, argmax, linparams):
gamma0 = self.space.axes[0].res # estimated FWHM on 10 pixels
gamma1 = self.space.axes[1].res
self.guess = [maximum, argmax[0], argmax[1], gamma0, gamma1, 0]
[docs]class PolarLorentzian2D(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x, y) = grid
(I, loc0, loc1, gamma0, gamma1, th, slope1, slope2, offset) = params
a, b = tuple(grid - center for grid, center in zip(rot2d(x, y, th), rot2d(loc0, loc1, th)))
return (I / (1 + (a / gamma0)**2 + (b / gamma1)**2) + x * slope1 + y * slope2 + offset)
[docs] def set_guess(self, maximum, argmax, linparams):
gamma0 = self.space.axes[0].res # estimated FWHM on 10 pixels
gamma1 = self.space.axes[1].res
self.guess = [maximum, argmax[0], argmax[1], gamma0, gamma1, 0, linparams[0], linparams[1], linparams[2]]
[docs] def integrate_signal(self):
return self.func(self.cxdata, (self.result[0], self.result[1], self.result[2], self.result[3], self.result[4], self.result[5], 0, 0, 0)).sum()
[docs]class Lorentzian2D(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x, y) = grid
(I, loc0, loc1, gamma0, gamma1, th, slope1, slope2, offset) = params
a, b = tuple(grid - center for grid, center in zip(rot2d(x, y, th), rot2d(loc0, loc1, th)))
return (I / (1 + (a/gamma0)**2) * 1 / (1 + (b/gamma1)**2) + x * slope1 + y * slope2 + offset)
[docs] def set_guess(self, maximum, argmax, linparams):
gamma0 = 5 * self.space.axes[0].res # estimated FWHM on 10 pixels
gamma1 = 5 * self.space.axes[1].res
self.guess = [maximum, argmax[0], argmax[1], gamma0, gamma1, 0, linparams[0], linparams[1], linparams[2]]
[docs]class Lorentzian2Dnobkg(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x, y) = grid
(I, loc0, loc1, gamma0, gamma1, th) = params
a, b = tuple(grid - center for grid, center in zip(rot2d(x, y, th), rot2d(loc0, loc1, th)))
return (I / (1 + (a/gamma0)**2) * 1 / (1 + (b/gamma1)**2))
[docs] def set_guess(self, maximum, argmax, linparams):
gamma0 = 5 * self.space.axes[0].res # estimated FWHM on 10 pixels
gamma1 = 5 * self.space.axes[1].res
self.guess = [maximum, argmax[0], argmax[1], gamma0, gamma1, 0]
[docs]class Lorentzian(AutoDimensionFit):
dimensions = {1: Lorentzian1D, 2: PolarLorentzian2D}
[docs]class Gaussian1D(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x,) = grid
(loc, I, sigma, offset, slope) = params
return I * numpy.exp(-((x-loc)/sigma)**2/2) + offset + x * slope
[docs]class Voigt1D(PeakFitBase):
[docs] @staticmethod
def func(grid, params):
(x, ) = grid
(I, loc, sigma, gamma, slope, offset) = params
z = (x - loc + numpy.complex(0, gamma)) / (sigma * numpy.sqrt(2))
return I * numpy.real(scipy.special.wofz(z))/(sigma * numpy.sqrt(2 * numpy.pi)) + offset + x * slope
[docs] def set_guess(self, maximum, argmax, linparams):
gamma0 = 5 * self.space.axes[0].res # estimated FWHM on 10 pixels
self.guess = [maximum, argmax[0], 0.01, gamma0, linparams[0], linparams[1]]