ccfcenter.py

Go to the documentation of this file.

#!/usr/bin/python
##
# 
# This module finds the center and width of a MARVELS ccf
# 
# This module uses a non-linear least-squares optimizer from scipy to 
#     fit a Gaussian to the central pixels of a ccf file.
# 
# @package ccfcenter
# @author Nathan De Lee
# @version $Revision: 1.2 $
# @date $Date: 2012-01-13 20:12:24 $
# 

# -----------------------------
# Standard library dependencies
# -----------------------------
import os
import sys
import collections

# -------------------
# Third-party imports
# -------------------
from scipy import optimize
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.backends.backend_pdf import PdfPages

##
# 
#     Choose which section of the ccf to fit based on width parameter 
#         or automatically using percent
#     
#     @param pixel   Numpy array of pixel positions
#     @param ccf     Numpy array of ccf values
#     @param width   Integer number of pixels to include, if None: 
#                     then automatically determine.
#     @param percent Floating point percentage of the singal amplitude 
#                     above the noise that a point must be to be included in 
#                     the fit.
#     
#     @return This function returns a tuple (pixel,ccf,lmargin,rmargin
#             ,noiselevel).
#     
#     @retval pixel      Numpy array of pixel positions
#     @retval ccf        Numpy array of ccf values
#     @retval lmargin    Integer index of the left margin of the input arrays
#     @retval rmargin    Integer index of the right margin of the input arrays
#     @retval noiselevel Floating point values of the CCF plot noise level
#     
def cutccf(pixel,ccf,width=None,percent=0.30): 
    cent_index= ccf.argmax()
    if width == None:
        noiselevel = np.median(ccf)
        cutoff = (percent * (max(ccf) - noiselevel))+noiselevel
        indice = np.array((ccf < cutoff).nonzero()).flatten()
        lmargin = max(indice[np.array((indice < cent_index).nonzero()).flatten()])
        rmargin = min(indice[np.array((indice > cent_index).nonzero()).flatten()])    
    else:
        lmargin = cent_index-width
        rmargin = cent_index+width+1
         
    return(pixel[lmargin:rmargin],ccf[lmargin:rmargin],lmargin,rmargin,noiselevel)

##
# 
#     Takes x points and returns a Gaussian function
#     
#     @param P A 3 element array (Amplitude,Offset,Sigma)
#     @param x Numpy array of x positions
#     
#     @returns y values for the x positions and parameters given
#     
#     @retval y Numpy array of y values 
#     
def gaussfunc(P,x):
        
    ##  @page Formula Mathematical Formula
    #   @section gauss Gaussian 
    #    @{
    #   I use this formulation for my Gaussian Equation
    #   @f$ y = Amp \exp(-\frac{(x - Offset)^{2}}{2\sigma^{2}})@f$
    ##   @}

    return(P[0]*np.exp(-(x-P[1])**2/(2*P[2]**2)))
  
##
# 
#     Creates a Gaussian signal and then fits it to see if I can recover the input parameters.
#     
#     @return None
#     
#     
def gausstest():
    ntrials = 10000
    #Open output pdf
    pp = PdfPages('test.pdf')  
    print "Progress: ",
    for i in range(ntrials) :
        if ntrials >= 10 and i % (ntrials/10) == 0 :
            print str("."),
        (x,y) = make_gauss(6,amp=0.91,sigma=10.4,offset=5.1,error=0.015)
        #Fit a gaussian to the data cutout     
        #(params,errors,chi)=fitgauss(x,y,inguess=[1,1,20],inerror=(np.zeros(len(x))+.015))
        (params,errors,chi)=fitgauss(x,y,inguess=[1,1,20])

        #print i,params,errors,chi
        if i == 0:
            allparams = params
            allerrors = errors
            allchi = chi
        else:
            allparams = np.vstack([allparams,params])
            allerrors = np.vstack([allerrors,errors])
            allchi = np.vstack([allchi,chi])

    ##@fn def gausstest():
    #@todo Make a histogram of all the parameters
    
    #Print out how well I did.
    print "\nActual:  ",np.std(allparams[:,0]),np.std(allparams[:,1]),np.std(allparams[:,2]) 
    print "First:   ",allerrors[0,0],allerrors[0,1],allerrors[0,2]
    print "Average: ",np.mean(allerrors[:,0]),np.mean(allerrors[:,1]),np.mean(allerrors[:,2])
    pp.close()
    return(0)

##
# 
#     Use a Levenberg-Marquadt non-linear least-squares fitter 
#     to fit a Gaussian to Xin and Yin
#     
#     The scipy module optimize is used to find the smallest residual
#     to @ref error which is compared to the @ref gauss function.
#     
#     @param xin Numpy array of x-values
#     @param yin Numpy array of y-values
#     @param inerror Numpy array of error values
#     @param inguess A 3 element list (Amplitude,Offset,Sigma) of 
#             starting values for Gaussian parameters 
#     
#     @return This function returns a tuple (Params, perror , red_chi)
#     
#     @retval Params A 3 element Numpy array of Gaussian 
#              Parameters (Amplitude,Offset,Sigma)
#     @retval perror A 3 element Numpy array of parameter errors 
#              calculated from the covariance matrix
#     @retval red_chi A float giving the reduced @f$\chi^{2}@f$ 
#              of the fit.
#    
def fitgauss(xin,yin,inerror=[],inguess=[1,1,1]):
    
    #Don't want my np.arrays messed up.
    x = np.array(xin,dtype='float64')
    y = np.array(yin,dtype='float64')
    guess = np.array(inguess[:])
    if len(inerror) == len(xin):
        error = np.array(inerror,dtype='float64')
    else:
        error = np.ones(len(xin)) 

    if(len(inerror) == len(xin)):
    ##  @page Formula Mathematical Formula
    #   @section error Error Function 
    #    @{
    #   I use this formulation for optimizing my least-squares fitter
    #    - With Errors:    @f$\frac{y - Gaussian(P,x)}{errror}@f$
    #    - Without Errors: @f$y - Gaussian(P,x)@f$
    #   @}
        errfunc = lambda P,x,y,error: (y - gaussfunc(P, x))/error 
        out = optimize.leastsq(errfunc, guess, args=(x, y, error),full_output=1)
    else:
        errfunc = lambda P,x,y: y - gaussfunc(P, x) 
        out = optimize.leastsq(errfunc, guess, args=(x, y),full_output=1)
 
    Params = out[0]
    #Sigma should be positive
    Params[2] = np.abs(Params[2])
    
    if out[1] != None :
        cov = out[1] 
        perror = np.array([np.sqrt(cov[0][0]),np.sqrt(cov[1][1]),np.sqrt(cov[2][2])])
    else:
        perror = np.array([1e11,1e11,1e11])
    if(len(inerror) == len(xin)):  
        red_chi = sum((errfunc(Params,x,y,error))**2)/(len(x)-3)
    else:
        red_chi = sum((errfunc(Params,x,y))**2)/(len(x)-3)
    
    ##@fn def fitgauss(xin,yin,inerror=[],inguess=[1,1,1]):
    #@remarks
    #If there are no measurement errors then the parameter errors must be scaled by the reduced chisq
    #Section 15.2 in <a href="http://adsabs.harvard.edu/abs/1992nrca.book.....P/">Numerical 
    #Recipes </a> discusses this.
    
    if(len(inerror) == 0):    
        perror = perror * np.sqrt(red_chi)
     
    return(Params,perror,red_chi)

##
# 
#     Generate Gaussian data using @ref gauss
# 
#     @param npts          Integer number of points to create
#     @param amp           Floating point amplitude of the Gaussian
#     @param offset        Floating point offset of the Gaussian
#     @param sigma         Floating point sigma of the Gaussian
#     @param error         Floating point number or a Numpy array of sigmas for the error in each point
#     @param inflate_error Floating point factor to scale the real as versus reported error
#     @param randpts       Bool True: Select x points from a uniform distribution;
#                          False Take from a regular grid.
#     
#     @return This function returns a tuple (x,y,y_err)
#     
#     @retval x     Numpy Array of x positions
#     @retval y     Numpy Array of y positions
#     @retval y_err Numpy Array of y errors
#    
def make_gauss(npts,amp=1,offset=0,sigma=1,error=0,inflate_error=1.0,randpts=False):
      
    if randpts == True:
        x = (np.random.uniform(size=npts) - 0.5)* 6.0 * sigma + offset
    else:
        x = (np.arange(npts,dtype='float') -npts/2.0) * 6 * sigma /npts + offset
    #Get Gaussian value
    y = gaussfunc([amp,offset,sigma],x)
    
    #Add Error to Gaussian
    if error > 0:
        y = y + (np.random.normal(0,error,npts)*inflate_error) 

    return(x,y)

##
# 
#     This function is the main() program for this module
#     
#     @param filename String name of MARVELS ccf file to process
#     @param width    Integer number of pixels to fit around peak; 
#                     None: Automatically choose
#     
#     @return The function returns 0 for success or an error message.
#     
#     @retval ret 0 for success or an error message 
#     
def ccfcenter(filename,width=None):
    #Numpy arrays are tricky.  To make an empty shapeless numpy array np.ndarray((0,))
    #I want a dictionary of dictionary of numpy arrays.
    #The collections defaultdict command allows one to do this.  It makes a dictionary typed as a callable
    #function.  For instance if you use list() it will allow you to use a .append on it.
    #You must define a factory which is callable, so I use the lambda to make np.ndarray((0,)) into a 
    #callable function like list()
    ccfdict = collections.defaultdict(lambda : collections.defaultdict(lambda : np.ndarray((0,))))

    #Read in data file
    try:
        inputf = open(filename)
    except IOError:
        print "File %s could not be opened!" % filename
        return(1)
    
    try:
        for line in inputf:
            #Ignore comments and blank lines
            if line.startswith("#") or line.isspace():
                continue
            line.strip()    
            aline = line.split()
            ccfdict[aline[0]]['pixel'] = np.append(ccfdict[aline[0]]['pixel'],float(aline[8]))
            ccfdict[aline[0]]['ccf'] = np.append(ccfdict[aline[0]]['ccf'],float(aline[9]))
            ccfdict[aline[0]]['rv'] = float(aline[1])
            ccfdict[aline[0]]['rverr'] = float(aline[2])
            ccfdict[aline[0]]['ccfrv'] = float(aline[3])
            ccfdict[aline[0]]['ccfrverr'] = float(aline[4])
    finally:
        
        inputf.close()
    
    #Get base filename
    basename = os.path.splitext(filename)[0]    
    #Open output pdf
    pp = PdfPages(basename+'.pdf')  
    #Open output file
    outputf = open(basename+'.out','w')
    #Loop through each epoch 
    i=0
    for key in sorted(ccfdict.keys()):
        i += 1
        #Make data cutout
        (pixcut,ccfcut,lmargin,rmargin,noiselvl) = cutccf(ccfdict[key]['pixel'],np.array(ccfdict[key]['ccf']),width)        
        print lmargin,rmargin,noiselvl
        #Fit a gaussian to the data cutout     
        (params,errors,chi)=fitgauss(pixcut,ccfcut,inguess=[1,1,4])
        #Print the results
        print i,key,params,errors,chi
        rv = ccfdict[key]['rv']
        err = ccfdict[key]['rverr']
        rvccf = ccfdict[key]['ccfrv']
        ccferr = ccfdict[key]['ccfrverr']
        #Write out to file
        out = map(str,(i,key,rv,err,rvccf,ccferr,' '.join(map(str,params)),' '.join(map(str,errors)),chi))
        outputf.write(' '.join(out))
        outputf.write('\n')
        #Plot the results
        pl.xlabel("Pixel")
        pl.ylabel("CCF")
        pl.ylim(-.2,1.1)
        pl.title("Date: %s" % key)
        pl.plot(ccfdict[key]['pixel'][lmargin-20:rmargin+20],ccfdict[key]['ccf'][lmargin-20:rmargin+20],'go')
        pl.plot(pixcut,ccfcut,'ro')
        pixcut = np.array(pixcut,dtype='float64')
        modelx = np.arange(100) * ((pixcut.max() - pixcut.min()+20)/ 100.0) + pixcut.min()-10
        pl.plot(modelx,gaussfunc(params,modelx))
        pp.savefig()
        pl.clf()      
    pp.close()
    outputf.close()
    return(0)
if __name__ == '__main__':
    #Check to make sure we have at least one argument and no more than 2 arguments
    if len(sys.argv) < 2 or len(sys.argv) > 3:
        print 'Usage: ccf ccfilename {width}\n'
        print 'If test is given for the filename, then will run a montecarlo simulation'
        print 'ccfilename              Output filename from ccfsearch.pro'
        print 'width                   Optional width parameter (in pixels)\n'
        sys.exit(1)
    filename = sys.argv[1]
    if filename == 'test':
        gausstest()
        sys.exit(0)
    else:
        width = None
        ret = ccfcenter(filename,width)
        sys.exit(ret)

##
#@mainpage A Guassian CCF Fitter 
#This module is designed to fit a gaussian to the peak of a cross-correlation function.  
#In it's current form, it uses a non-linear least squares fitter from the scipy optimize module.
#@verbatim
#Usage: ccf ccfilename {width}
#If test is given for the filename, then will run a montecarlo simulation
#ccfilename              Output filename from ccfsearch.pro
#width                   Optional width parameter (in pixels)
#@endverbatim