Source code for aces.beamset.airyfitter

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
===========
airyfitter
===========
.. codeauthor:: Dvid McConnell

Derived from gaussfitter by Adam Ginsburg <adam.g.ginsburg@gmail.com> 3/17/08
Latest version available at <https://github.com/keflavich/gaussfitter>, where
it was moved from google code on January 30, 2014

"""

import numpy as np
from numpy.ma import median
from numpy import pi
import scipy.special as ssp
from mpfit.mpfit import mpfit


[docs]def moments(data, circle, rotate, vheight, estimator=median, angle_guess=45.0, **kwargs): """ Returns (height, amplitude, x, y, width_x, width_y, rotation angle) the parameters of a 2D airy pattern by calculating its moments. Depending on the input parameters, will only output a subset of the above. If using masked arrays, pass estimator=np.ma.median """ total = np.abs(data).sum() Y, X = np.indices(data.shape) # python convention: reverse x,y np.indices y = np.argmax((X*np.abs(data)).sum(axis=1)/total) x = np.argmax((Y*np.abs(data)).sum(axis=0)/total) col = data[int(y), :] # FIRST moment, not second! # width_x = np.sqrt(np.abs((np.arange(col.size)-y)*col).sum() / np.abs(col).sum()) # row = data[:, int(x)] # width_y = np.sqrt(np.abs((np.arange(row.size)-x)*row).sum() / np.abs(row).sum()) width_x = np.abs((np.arange(col.size)-y)*col).sum() / np.abs(col).sum() row = data[:, int(x)] width_y = np.abs((np.arange(row.size)-x)*row).sum() / np.abs(row).sum() width = (width_x + width_y) / 2. height = estimator(data.ravel()) amplitude = data.max()-height mylist = [amplitude, x, y] if np.isnan((width_y, width_x, height, amplitude)).any(): raise ValueError("something is nan") if vheight: mylist = [height] + mylist if not circle: mylist = mylist + [width_x, width_y] if rotate: # rotation "moment" is a little above zero to initiate the fitter # with something not locked at the edge of parameter space mylist = mylist + [angle_guess] # also, circles don't rotate. else: mylist = mylist + [width] return mylist
[docs]def airy(x): # Return airy function value at x # x in units of lambda/D return 2 * ssp.jn(1, x*pi) / x / pi
[docs]def airy_sq(x): # Return airy function value at x # x in units of lambda/D # b = 1.0/10.14 b = 0.22 ai1 = airy(x) ai2 = b * airy(x * b) ret = (ai1 - ai2)**2 return ret
[docs]def twodairy(inpars, circle=False, rotate=True, vheight=True, shape=None): """ Returns a 2d airy pattern of the form: x' = np.cos(rota) * x - np.sin(rota) * y y' = np.sin(rota) * x + np.cos(rota) * y (rota should be in degrees) g = b + a * airy (sqrt((x' - xc)/xw)**2 + ((y' - yc)/yw)**2 )) inpars = [b,a,center_x,center_y,width_x,width_y,rota] (b is background height, a is peak amplitude) where x and y are the input parameters of the returned function, and all other parameters are specified by this function However, the above values are passed by list. The list should be: inpars = (height,amplitude,center_x,center_y,width_x,width_y,rota) You can choose to ignore / neglect some of the above input parameters using the following options: Parameters ---------- circle : bool default is an elliptical airy (different x, y widths), but can reduce the input by one parameter if it's a circular airy pattern rotate : bool default allows rotation of the pattern ellipse. Can remove last parameter by setting rotate=0 vheight : bool default allows a variable height-above-zero, i.e. an additive constant for the function. Can remove first parameter by setting this to 0 shape : tuple if shape is set (to a 2-parameter list) then returns an image with the airy pattern defined by inpars """ """ Changes 2017 March: 1. pop parameters in the order amplitude, center_x, center_y; not amplitude, center_y, center_x 2. rota is applied to grid before function calculation; we want to negate it 3. rotgauss is called with y first: uses x as y, and y as x; reverse the argument order """ inpars_old = inpars inpars = list(inpars) if vheight: height = inpars.pop(0) height = float(height) else: height = float(0) amplitude, center_x, center_y = inpars.pop(0), inpars.pop(0), inpars.pop(0) amplitude = float(amplitude) center_x = float(center_x) center_y = float(center_y) if circle: width = inpars.pop(0) width_x = float(width) width_y = float(width) rotate = 0 else: width_x, width_y = inpars.pop(0), inpars.pop(0) width_x = float(width_x) width_y = float(width_y) if rotate: rota = inpars.pop(0) # rota = -pi/180. * float(rota) rota = -1.0 * float(rota) rcen_x = center_x * np.cos(rota) - center_y * np.sin(rota) rcen_y = center_x * np.sin(rota) + center_y * np.cos(rota) else: rcen_x = center_x rcen_y = center_y if len(inpars) > 0: raise ValueError("There are still input parameters:" + str(inpars) + " and you've input: " + str(inpars_old) + " circle=%d, rotate=%d, vheight=%d" % (circle, rotate, vheight)) def rotairy(y, x): if rotate: xp = x * np.cos(rota) - y * np.sin(rota) yp = x * np.sin(rota) + y * np.cos(rota) else: xp = x yp = y g = height + amplitude * airy_sq(np.sqrt(((rcen_x-xp)/width_x)**2 + ((rcen_y-yp)/width_y)**2)) return g if shape is not None: return rotairy(*np.indices(shape)) else: return rotairy
[docs]def airyfit(data, err=None, params=(), autoderiv=True, return_error=False, circle=False, fixed=np.repeat(False, 7), limitedmin=[False, False, False, False, True, True, True], limitedmax=[False, False, False, False, False, False, True], usemoment=np.array([], dtype='bool'), minpars=np.repeat(0, 7), maxpars=[0, 0, 0, 0, 0, 0, 180], rotate=True, vheight=True, quiet=True, returnmp=False, returnfitimage=False, **kwargs): """ Airy pattern fitter with the ability to fit a variety of different forms of 2-dimensional airy pattern. Parameters ---------- data : `numpy.ndarray` 2-dimensional data array err : `numpy.ndarray` or None error array with same size as data array. Defaults to 1 everywhere. params : (height, amplitude, x, y, width_x, width_y, rota) Initial input parameters for airy pattern. If not input, these will be determined from the moments of the system, assuming no rotation autoderiv : bool Use the autoderiv provided in the lmder.f function (the alternative is to us an analytic derivative with lmdif.f: this method is less robust) return_error : bool Default is to return only the airy parameters. If ``True``, return fit params & fit error returnfitimage : bool returns (best fit params,best fit image) returnmp : bool returns the full mpfit struct circle : bool The default is to fit an elliptical airy pattern (different x, y widths), but the input is reduced by one parameter if it's a circular pattern. rotate : bool Allow rotation of the ellipse. Can remove last parameter of input & fit by setting rotate=False. Angle should be specified in degrees. vheight : bool Allows a variable height-above-zero, i.e. an additive constant background for the function. Can remove the first fitter parameter by setting this to ``False`` usemoment : `numpy.ndarray`, dtype='bool' Array to choose which parameters to use a moment estimation for. Other parameters will be taken from params. Returns ------- (params, [parerr], [fitimage]) | (mpfit, [fitimage]) parameters : list The default output is a set of parameters with the same shape as the input parameters fitimage : `numpy.ndarray` If returnfitimage==True, the last return will be a 2D array holding the best-fit model mpfit : `mpfit` object If ``returnmp==True`` returns a `mpfit` object. This object contains a `covar` attribute which is the 7x7 covariance array generated by the mpfit class in the `mpfit_custom.py` module. It contains a `param` attribute that contains a list of the best fit parameters in the same order as the optional input parameter `params`. """ data = data.view(np.ma.MaskedArray) usemoment = np.array(usemoment, dtype='bool') params = np.array(params, dtype='float') if usemoment.any() and len(params) == len(usemoment): moment = np.array(moments(data, circle, rotate, vheight, **kwargs), dtype='float') params[usemoment] = moment[usemoment] elif params == [] or len(params) == 0: params = (moments(data, circle, rotate, vheight, **kwargs)) if not vheight: # If vheight is not set, we set it for sub-function calls but fix the # parameter at zero vheight = True params = np.concatenate([[0], params]) fixed[0] = 1 # mpfit will fail if it is given a start parameter outside the allowed range: for i, par in enumerate(params): if par > maxpars[i] and limitedmax[i]: params[i] = maxpars[i] if par < minpars[i] and limitedmin[i]: params[i] = minpars[i] # One time: check if error is set, otherwise fix it at 1. err = err if err is not None else 1.0 def mpfitfun(data, err): def f(p, fjac): twoda = twodairy(p, circle, rotate, vheight) delta = (data - twoda(*np.indices(data.shape))) / err return [0, delta.compressed()] return f parinfo = [{'n': 1, 'value': params[1], 'limits': [minpars[1], maxpars[1]], 'limited': [limitedmin[1], limitedmax[1]], 'fixed': fixed[1], 'parname': "AMPLITUDE", 'error': 0}, {'n': 2, 'value': params[2], 'limits': [minpars[2], maxpars[2]], 'limited': [limitedmin[2], limitedmax[2]], 'fixed': fixed[2], 'parname': "XSHIFT", 'error': 0}, {'n': 3, 'value': params[3], 'limits': [minpars[3], maxpars[3]], 'limited': [limitedmin[3], limitedmax[3]], 'fixed': fixed[3], 'parname': "YSHIFT", 'error': 0}, {'n': 4, 'value': params[4], 'limits': [minpars[4], maxpars[4]], 'limited': [limitedmin[4], limitedmax[4]], 'fixed': fixed[4], 'parname': "XWIDTH", 'error': 0}] if vheight: parinfo.insert(0, {'n': 0, 'value': params[0], 'limits': [minpars[0], maxpars[0]], 'limited': [limitedmin[0], limitedmax[0]], 'fixed': fixed[0], 'parname': "HEIGHT", 'error': 0}) if not circle: parinfo.append({'n': 5, 'value': params[5], 'limits': [minpars[5], maxpars[5]], 'limited': [limitedmin[5], limitedmax[5]], 'fixed': fixed[5], 'parname': "YWIDTH", 'error': 0}) if rotate: parinfo.append({'n': 6, 'value': params[6], 'limits': [minpars[6], maxpars[6]], 'limited': [limitedmin[6], limitedmax[6]], 'fixed': fixed[6], 'parname': "ROTATION", 'error': 0}) if not autoderiv: # the analytic derivative, while not terribly difficult, is less # efficient and useful. I only bothered putting it here because I was # instructed to do so for a class project - please ask if you would # like this feature implemented raise NotImplementedError("I'm sorry, I haven't implemented this feature yet. " "Given that I wrote this message in 2008, " "it will probably never be implemented.") else: mp = mpfit(mpfitfun(data, err), parinfo=parinfo, quiet=quiet) if mp.errmsg: raise Exception("MPFIT error: {0}".format(mp.errmsg)) if (not circle) and rotate: mp.params[-1] %= pi # mp.params[-1] %= 180.0 mp.chi2 = mp.fnorm try: mp.chi2n = mp.fnorm/mp.dof except ZeroDivisionError: mp.chi2n = np.nan if returnmp: returns = (mp) elif return_error: returns = mp.params, mp.perror else: returns = mp.params if returnfitimage: fitimage = twodairy(mp.params, circle, rotate, vheight)(*np.indices(data.shape)) returns = (returns, fitimage) return returns