1 # gaussfitter.py
   2 # created by Adam Ginsburg (adam.ginsburg@colorado.edu or keflavich@gmail.com) 3/17/08)
   3 from numpy import *
   4 from scipy import optimize
   5 from scipy import stats
   6 
   7 def moments(data,circle,rotate,vheight):
   8     """Returns (height, amplitude, x, y, width_x, width_y, rotation angle)
   9     the gaussian parameters of a 2D distribution by calculating its
  10     moments.  Depending on the input parameters, will only output 
  11     a subset of the above"""
  12     total = data.sum()
  13     X, Y = indices(data.shape)
  14     x = (X*data).sum()/total
  15     y = (Y*data).sum()/total
  16     col = data[:, int(y)]
  17     width_x = sqrt(abs((arange(col.size)-y)**2*col).sum()/col.sum())
  18     row = data[int(x), :]
  19     width_y = sqrt(abs((arange(row.size)-x)**2*row).sum()/row.sum())
  20     width = ( width_x + width_y ) / 2.
  21     height = stats.mode(data.ravel())[0][0]
  22     amplitude = data.max()-height
  23     mylist = [amplitude,x,y]
  24     if vheight==1:
  25         mylist = [height] + mylist
  26     if circle==0:
  27         mylist = mylist + [width_x,width_y]
  28     else:
  29         mylist = mylist + [width]
  30     if rotate==1:
  31         mylist = mylist + [0.] #rotation "moment" is just zero...
  32     return tuple(mylist)
  33 
  34 def twodgaussian(inpars, circle, rotate, vheight):
  35     """Returns a 2d gaussian function of the form:
  36         x' = cos(rota) * x - sin(rota) * y
  37         y' = sin(rota) * x + cos(rota) * y
  38         (rota should be in degrees)
  39         g = b + a exp ( - ( ((x-center_x)/width_x)**2 +
  40         ((y-center_y)/width_y)**2 ) / 2 )
  41 
  42         where x and y are the input parameters of the returned function,
  43         and all other parameters are specified by this function
  44 
  45         However, the above values are passed by list.  The list should be:
  46         inpars = (height,amplitude,center_x,center_y,width_x,width_y,rota)
  47 
  48         You can choose to ignore / neglect some of the above input parameters using the following options:
  49             circle=0 - default is an elliptical gaussian (different x, y widths), but can reduce
  50                 the input by one parameter if it's a circular gaussian
  51             rotate=1 - default allows rotation of the gaussian ellipse.  Can remove last parameter
  52                 by setting rotate=0
  53             vheight=1 - default allows a variable height-above-zero, i.e. an additive constant
  54                 for the Gaussian function.  Can remove first parameter by setting this to 0
  55         """
  56     inpars_old = inpars
  57     inpars = list(inpars)
  58     if vheight == 1:
  59         height = inpars.pop(0)
  60         height = float(height)
  61     else:
  62         height = float(0)
  63     amplitude, center_x, center_y = inpars.pop(0),inpars.pop(0),inpars.pop(0)
  64     amplitude = float(amplitude)
  65     center_x = float(center_x)
  66     center_y = float(center_y)
  67     if circle == 1:
  68         width = inpars.pop(0)
  69         width_x = float(width)
  70         width_y = float(width)
  71     else:
  72         width_x, width_y = inpars.pop(0),inpars.pop(0)
  73         width_x = float(width_x)
  74         width_y = float(width_y)
  75     if rotate == 1:
  76         rota = inpars.pop(0)
  77         rota = pi/180. * float(rota)
  78         rcen_x = center_x * cos(rota) - center_y * sin(rota)
  79         rcen_y = center_x * sin(rota) + center_y * cos(rota)
  80     else:
  81         rcen_x = center_x
  82         rcen_y = center_y
  83     if len(inpars) > 0:
  84         raise ValueError("There are still input parameters:" + str(inpars) + \
  85                 " and you've input: " + str(inpars_old) + " circle=%d, rotate=%d, vheight=%d" % (circle,rotate,vheight) )
  86             
  87     def rotgauss(x,y):
  88         if rotate==1:
  89             xp = x * cos(rota) - y * sin(rota)
  90             yp = x * sin(rota) + y * cos(rota)
  91         else:
  92             xp = x
  93             yp = y
  94         g = height+amplitude*exp(
  95             -(((rcen_x-xp)/width_x)**2+
  96             ((rcen_y-yp)/width_y)**2)/2.)
  97         return g
  98     return rotgauss
  99 
 100 def gaussfit(data,err=None,params=[],autoderiv=1,return_all=0,circle=0,rotate=1,vheight=1):
 101     """
 102     Gaussian fitter with the ability to fit a variety of different forms of 2-dimensional gaussian.
 103     
 104     Input Parameters:
 105         data - 2-dimensional data array
 106         err=None - error array with same size as data array
 107         params=[] - initial input parameters for Gaussian function.
 108             (height, amplitude, x, y, width_x, width_y, rota)
 109             if not input, these will be determined from the moments of the system, 
 110             assuming no rotation
 111         autoderiv=1 - use the autoderiv provided in the lmder.f function (the alternative
 112             is to us an analytic derivative with lmdif.f: this method is less robust)
 113         return_all=0 - Default is to return only the Gaussian parameters.  See below for
 114             detail on output
 115         circle=0 - default is an elliptical gaussian (different x, y widths), but can reduce
 116             the input by one parameter if it's a circular gaussian
 117         rotate=1 - default allows rotation of the gaussian ellipse.  Can remove last parameter
 118             by setting rotate=0
 119         vheight=1 - default allows a variable height-above-zero, i.e. an additive constant
 120             for the Gaussian function.  Can remove first parameter by setting this to 0
 121 
 122     Output:
 123         Default output is a set of Gaussian parameters with the same shape as the input parameters
 124         Can also output the covariance matrix, 'infodict' that contains a lot more detail about
 125             the fit (see scipy.optimize.leastsq), and a message from leastsq telling what the exit
 126             status of the fitting routine was
 127 
 128         Warning: Does NOT necessarily output a rotation angle between 0 and 360 degrees.
 129     """
 130     if params == []:
 131         params = (moments(data,circle,rotate,vheight))
 132     if err == None:
 133         errorfunction = lambda p: ravel((twodgaussian(p,circle,rotate,vheight)(*indices(data.shape)) - data))
 134     else:
 135         errorfunction = lambda p: ravel((twodgaussian(p,circle,rotate,vheight)(*indices(data.shape)) - data)/err)
 136     if autoderiv == 0:
 137         # the analytic derivative, while not terribly difficult, is less efficient and useful.  I only bothered
 138         # putting it here because I was instructed to do so for a class project - please ask if you would like 
 139         # this feature implemented
 140         raise ValueError("I'm sorry, I haven't implemented this feature yet.")
 141     else:
 142         p, cov, infodict, errmsg, success = optimize.leastsq(errorfunction, params, full_output=1)
 143     if  return_all == 0:
 144         return p
 145     elif return_all == 1:
 146         return p,cov,infodict,errmsg