Fitting.py

'''
Fitting.py: ERE created 12/27/2013
This file will contain the parts and pieces to do fitting 
in the analysis window (and possibly in the measurement window eventually))
'''

'''
Some useful things I have figured out so far
I am going to fit using lmfit package
My fitting functions will be defined in the "FitFunctions" package
Each function will get a single file that must define:
paramList : A simple list of the parameter keys (ideally the parameter keys are also defined as variables)
Residuals : a residuals function that takes in (paramsObj, xData, yData,*args,**kwargs)
            and returns Fcn(params,xData) - yData (possibly eventually with some error adjustment)
Function : The actual function to be fit that takes in (paramsObj,xData,*args,**kwargs)

lmfit:
    Parameters() creates a new empty parameters object
    parameters can be added to the object using params.add('key',value=XXX,min=XXX,max=XXX,vary=T/F,expr=XXX)
    add_many also exists, but is a little more finicky
    
    Fit is performed by: result = lmfit.minimize(ResidualsFcn,paramsObj,args = (xdata,ydata,*args, **kwargs))
    
    lmfit.report_fit(paramsObj) will print out the basic details that include parameter values with error
    And will give correlations.
    It is also possible to do something with result and directly with paramsObj.
    Default returns standard error, can also call conf_interval to get the actual intervals - does math

To get the list of functions in the FitFunctions Package:
import pkgutil
for importer,modname,ispkg in pkgutil.iter_modules(FitFunctions.__path__): 
    print modname,ispkg 
modname is really what I want, it's a string with the name of the module (fitFunction name)
to load/access a function by string: use importlib.import_module?

Consider storing a default starting parameters config in the function file
'''

# I plan to have this module generate a dock that can be inserted into a pyqtgraph.dockarea,
# Or should I have this one generate the widget, and insert the widget into the dock in the calling class?
from PyQt4 import QtGui, QtCore
#import FitFunctions
import lmfit

StartKey = 'Start'
MinKey = 'Min'
MaxKey = 'Max'
VaryKey = 'Vary'
ResultKey = 'Result'
ErrorKey = 'Std Err'
ConfIntKey = 'Conf Int'
CorrelKey = 'Correlations'

tableHeaders = [StartKey,MinKey,MaxKey,VaryKey,ResultKey,ErrorKey,ConfIntKey,CorrelKey]

def FitOne(functionName,params,xData,yData):
    import lmfit
    fcn = GetFunctionFromName(functionName)
    result = lmfit.minimize(fcn.Residuals,params,args=(xData,yData))
    return result
    
def GetFunctionFromName(functionName):
    modToImport = 'FitFunctions.'+functionName
    # now convert to the module for that function
    currFunction = __import__(modToImport,fromlist=[''])
    return currFunction

def GetParameterValue(parameterObj,paramName):
    import Fitting
    if paramName == Fitting.StartKey:
        return parameterObj.init_value
    elif paramName == Fitting.MinKey:
        return parameterObj.min
    elif paramName == Fitting.MaxKey:
        return parameterObj.max
    elif paramName == Fitting.VaryKey:
        return parameterObj.vary
    elif paramName == Fitting.ResultKey:
        if parameterObj.value == parameterObj.init_value:
            return None
        else:
            return parameterObj.value
    elif paramName == Fitting.ErrorKey:
        return parameterObj.stderr
    elif paramName == Fitting.ConfIntKey:
        return None
    elif paramName == Fitting.CorrelKey:
        return parameterObj.correl

class SingleFit(QtGui.QWidget):
    
    def __init__(self,parent=None,ROIs = None):
        super(SingleFit,self).__init__(parent)
        from AutoGeneratedFiles.FitSingleCurveWidget import Ui_FitSingleCurveWidget
        self.ui = Ui_FitSingleCurveWidget()
        self.ui.setupUi(self)
        

        self.parent = parent
        self.lastFilename = ''
        
        self.currFuncName = None
        # make a blank parameters object
        self.params = lmfit.Parameters()
        # Now populate the drop downs
        self.ROIs = ROIs
        # generate ROIlist
        self.ROIlist = []
        for ROI in self.ROIs:
            self.ROIlist.append(ROI.name)
        # set up the ROI list
        self.ui.comboROI.addItems(self.ROIlist)
        
        # then populate the Fit Function list
        self.functionList = []
        import pkgutil
        for importer,modname,ispkg in pkgutil.iter_modules(FitFunctions.__path__): 
            self.functionList.append(modname)
        self.ui.comboFitFunction.addItems(self.functionList)

        # populate the new function (will also update the parameter table)
        self.updateNewFunction()
        
        # connect signals and slots
        self.ui.comboFitFunction.currentIndexChanged.connect(self.updateNewFunction)
        # Set up a link with the change in the function to update the table with a new param object
        
        self.ui.pushDetect.clicked.connect(self.detectStartingParams)
        self.ui.pushFit.clicked.connect(self.fit)
        self.ui.pushPlot.clicked.connect(self.updatePlot)
        self.ui.pushSave.clicked.connect(self.saveResults)
    
    def detectStartingParams(self):
        print 'Eventually I will autodetect reasonable starting parameters'
        import Fitting
        xData,yData = self.getCurrentData()
        currFunc = self.getCurrentFunction()
        # update the parameters to save any user entered information first
        self.setParametersFromTable()
        currFunc.AutoDetectValues(self.params,xData,yData)
        # now update the table
        self.updateParametersTable()
    
    def fit(self):
        # set everything up, then call the fit function, finally print out the results in the table
        self.setParametersFromTable()
        import Fitting
        xData,yData = self.getCurrentData()
        result = Fitting.FitOne(self.getCurrentFunctionName(),self.params,xData,yData)
        import lmfit
        lmfit.report_fit(self.params)
        self.updateParametersTable()
        self.updatePlot()
    
    def getCurrentData(self):
        # read the ROI box, get the ROI I need
        ROIname = str(self.ui.comboROI.currentText())
        for ROI in self.ROIs:
            if ROIname == ROI.name:
                # Found it!
                return ROI.xData,ROI.yData
        else:
            return None, None

    def getCurrentROI(self):
        # read the ROI box, get the ROI I need
        ROIname = str(self.ui.comboROI.currentText())
        for ROI in self.ROIs:
            if ROIname == ROI.name:
                # Found it!
                return ROI
        else:
            return None
        
    def getCurrentFunction(self):
        currFunctionName = self.getCurrentFunctionName()
        import Fitting
        return Fitting.GetFunctionFromName(currFunctionName)
    
    def getCurrentFunctionName(self):
        return self.currFuncName
    
    def saveResults(self):
        # save the current table to a file
        import os.path
        if self.parent is not None:
            dirString = os.path.join(self.parent.defaultDirectory,self.lastFilename)
            if self.parent.lastDirectory != '':
            # override default, use last
                dirString = os.path.join(self.parent.lastDirectory,self.lastFilename)
        fname = str(QtGui.QFileDialog.getSaveFileName(directory = dirString,filter = '*.txt;;*.dat;;*.csv;;*.*', selectedFilter = '*.dat'))#set the default filetype to save. 
        
        if not fname:
            return
        else:
            oldDir,filename = os.path.split(fname)
            if oldDir != '' and self.parent is not None:
                self.parent.lastDirectory = oldDir
            if filename != '':
                self.lastFilename = filename
                
            # now actually save the table to file
            import csv
            with open(fname, 'wb') as csvfile:
                writer = csv.writer(csvfile, delimiter='\t')
                # Write the header line - initial 'Parameter' gives space for a header label in each row
                headerLine = ['Parameter']
                for colNum in xrange(self.ui.tableParameters.columnCount()):
                    # read the column name and check against values I need for fitting - maybe put in to a function?
                    colName = str(self.ui.tableParameters.horizontalHeaderItem(colNum).text())
                    headerLine.append(colName)
                writer.writerow(headerLine)
                # cycle through the whole table
                for rowNum in xrange(self.ui.tableParameters.rowCount()):
                    line = []
                    # start with the parameter name
                    line.append(str(self.ui.tableParameters.verticalHeaderItem(rowNum).text()))
                    for colNum in xrange(self.ui.tableParameters.columnCount()):
                        try:
                            value = str(self.ui.tableParameters.item(rowNum,colNum).text())
                        except AttributeError:
                            # I hit a blank line that does not have a text value
                            value = ''
                        line.append(value)
                    writer.writerow(line)
    
    def setParametersFromTable(self):
        # DO THIS WITH A DICTIONARY!!!! Keys are in the tableHeader or row header
        import Fitting
        from MeasureUtils import strToBool
        for rowNum in xrange(self.ui.tableParameters.rowCount()):
            value = None
            minVal = None
            maxVal = None
            varyVal = True
            paramName = str(self.ui.tableParameters.verticalHeaderItem(rowNum).text())
            for colNum in xrange(self.ui.tableParameters.columnCount()):
                # read the column name and check against values I need for fitting - maybe put in to a function?
                colName = str(self.ui.tableParameters.horizontalHeaderItem(colNum).text())
                try:
                    textVal = self.ui.tableParameters.item(rowNum,colNum).text()
                except AttributeError:
                    # current cell is empty, keep going
                    continue
                if colName == Fitting.StartKey:
                    try:
                        value = float(textVal)
                    except (ValueError,AttributeError):
                        print 'Error converting to float',paramName,colName
                        value = None
                elif colName == Fitting.MinKey:
                    try:
                        minVal = float(textVal)
                    except (ValueError,AttributeError):
                        print 'Error converting to float',paramName,colName
                        minVal = None
                elif colName == Fitting.MaxKey:
                    try:
                        maxVal = float(textVal)
                    except (ValueError,AttributeError):
                        print 'Error converting to float',paramName,colName
                        maxVal = None
                elif colName == Fitting.VaryKey:
                    try:
                        varyVal = strToBool(textVal)
                    except (ValueError,AttributeError):
                        print 'Error converting to boolean',paramName,colName
                        varyVal = True
            self.params.add(paramName,value = value,min = minVal, max=maxVal,vary = varyVal)
    
    def updateNewFunction(self):
        # Make sure the function changed
        newFuncName = str(self.ui.comboFitFunction.currentText())
        if newFuncName is self.currFuncName:
            return
        
        self.currFuncName = newFuncName
        # I just changed the function in the drop down, create a new params object
        self.params = lmfit.Parameters()
        # set any default values here
        # clear the table
        self.ui.tableParameters.clearContents()
        # Then populate the table
        self.updateParametersTable()
    
    def updateParametersTable(self):
        import Fitting
        # read the currently selected function
        currFunction = self.getCurrentFunction()
        # set the table size
        self.ui.tableParameters.setColumnCount(len(Fitting.tableHeaders))
        self.ui.tableParameters.setRowCount(len(currFunction.paramList))
        self.ui.tableParameters.setVerticalHeaderLabels(currFunction.paramList)
        self.ui.tableParameters.setHorizontalHeaderLabels(Fitting.tableHeaders)
        
        # set any defaults? - Handle in updatNewFunction
        
        # read any parameters results (if any)
        # Use Param.stderr != None to determine if a fit has been run or just read it all and set blank for None
        from itertools import izip
        for colNum, colName in izip(xrange(len(Fitting.tableHeaders)),Fitting.tableHeaders):
            for rowNum,rowName in izip(xrange(len(currFunction.paramList)),currFunction.paramList):
                #I will now hit each cell in the table
                try:
                    currParam = self.params[rowName]
                except KeyError:
                    continue
                value = Fitting.GetParameterValue(currParam,colName)
                # convert the value to a string appropriately
                if value == None or value==float('-inf') or value == float('inf'):
                    value = ''
                elif isinstance(value,bool):
                    value = '%s'%value
                else:
                    try:
                        # I'm a number
                        value = '%e'%value
                    except TypeError:
                        # I'm not a number, could be boolean or string or dict
                        value = '%s'%value
                self.ui.tableParameters.setItem(rowNum,colNum,QtGui.QTableWidgetItem(value))
    
    def updatePlot(self):
        #update the plot of the current ROI to give it an overlay
        ROI = self.getCurrentROI()
        xdata,ydata = self.getCurrentData()
        currFunction = self.getCurrentFunction()
        modelData = currFunction.Function(self.params,xdata)
        ROI.plot(xdata,modelData,pen='g')