Source code for binoculars.fit

import numpy
import scipy.optimize
import scipy.special
import inspect
import re


[docs]class FitBase: 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(r"\((.*?)\)", line)[0].split(",") for line in code.split("\n")[2:4] ) if space.dimension != len(args[0]): raise ValueError( f"dimension mismatch: space has {space.dimension}," f" {self.__class__.__name__} expects {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( f"invalid number of guess parameters {guess!r}" f" for {self.parameters!r}" ) 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(r"\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( f"{n}: {p:.4g} +/- {v:.4g}" 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 \ f"{self.__class__.__name__} fit on" \ f" {self.space}\n{self.message}\n{self.summary}"
[docs]class PeakFitBase(FitBase): def __init__(self, space, guess=None, loc=None): if loc is not None: self.argmax = tuple(loc) else: self.argmax = None super().__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 is 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 is not 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, rcond=None)[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( f"{space.dimension}-dimensional space" f" not supported for {cls.__name__}" )
# 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(f"unsupported fit function '{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) = undefined # TODO 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]]