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