fitting_functions.py

import os, copy
import numpy as N
import simplejson
from collections import deque

class FunctionParamGroup(object):
    def __init__(self, paramNames=None):
        self.params = {}
        
        if paramNames is None:
            paramNames = []
        for paramName in paramNames:
            self.params.update({ paramName: FunctionParam(paramName) })
    
    def get(self, paramName):
        return self.params[paramName].value
    
    def set(self, paramName, value):
        self.params[paramName].value = value
    
    def asArray(self):
        return self.params.values()
    
    def __repr__(self):
        return str(self.asArray())
        

class FunctionParam(object):
    def __init__(self, paramName = '', initValue=0):
        self.paramName = paramName
        self.value = initValue
        self.error = 0
        self.upperLimit = 0
        self.lowerLimit = 0
    
    def __repr__(self):
        return str(self.getJSON())
    
    def getJSON(self):
        return { 'paramName': self.paramName, 'value': self.value, 'error': self.error, 'upperLimit': self.upperLimit, 'lowerLimit': self.lowerLimit }
        
        

class FunctionGroup(object):
    def __init__(self):
        self.functions = []
        self.data = {}
        
    def __repr__(self):
        foo = "["
        for function in self.functions:
            foo += function.__repr__() + ", "
        return foo + "]"
    
    def foo(self):
        f1 = Linear()
        f2 = Linear()
        f1.functionParams.params = { 'X1': FunctionParam('X1', 0), 'X2': FunctionParam('X2', 1),
                                     'Y1': FunctionParam('Y1', 2), 'Y2': FunctionParam('Y2', 2) }
        f2.functionParams.params = { 'X1': FunctionParam('X1', 0), 'X2': FunctionParam('X2', 1),
                                     'Y1': FunctionParam('Y1', 3), 'Y2': FunctionParam('Y2', 3) }
        self.functions = [f1, f2]
    
    def getValueAtX(self, X):
        functionYs = []
        for function in self.functions:
            value = function.function(X, None)
            if value:
                functionYs.append(value)
        
        return sum(functionYs)
    
    def function(self, Domain, Range):
        functionRanges = []
        for function in self.functions:
            functionRanges.append(function.function(Domain, Range))
        
        functionRange = sum(N.array(functionRanges))
        return functionRange
    
    def chisq(self, xData, yData, yErr):
        yCalc = self.function(xData, yData)
        
        yErr_temp = copy.deepcopy(yErr)
    #    zero_loc = N.where(yErr == 0)[0]
    #    if len(zero_loc) != 0:
    #        yErr_temp[zero_loc] = 1.0
        chi = ((yData - yCalc) / yErr_temp) ** 2
        
        return chi.sum() / (len(yData) - self.countFunctionParams())
    
    def countFunctionParams(self):
        count = 0
        for function in self.functions:
            count += len(function.functionParams.params)
        return count
            

    def createFunction(self, xData, yData):
        functionDomain    = N.arange(min(xData), max(xData), abs(max(xData) - min(xData)) / 200.)
        functionDataRange = N.arange(min(yData), max(yData), abs(max(yData) - min(yData)) / 200.)

        print self

        functionRanges = []
        functionYs = []
        functionInfos = []
        for function in self.functions:
            functionRanges.append(function.function(functionDomain, functionDataRange))
            functionYs.append(function.function(xData, yData))
            functionInfos.append(function.getJSON())
        
        functionRange = sum(N.array(functionRanges))
        functionData = zip_(functionDomain, functionRange)
        
        functionY = sum(N.array(functionYs))
        functionResiduals = N.subtract(functionY, yData)
        functionResidualsScaled = N.divide(functionResiduals, self.countFunctionParams())
        #functionResidualData = zip_(xData, functionResiduals)
        functionResidualData = zip_(xData, functionResidualsScaled)

        print
        print functionRange
        print functionY
        print
        print functionData
        print functionResidualData
        print
        print '--------'
        print
        
        functionObj = dict(fit=functionData, resid=functionResidualData, functionInfos=functionInfos)
        return functionObj
    
    
    
    def defPoint(self, request, functionIndex=-1):
        xPos = float(request.POST['xPos'])
        yPos = float(request.POST['yPos'])
        
        # We need to subtract the value from rest of the functions
        if len(self.functions) > 1:
            yPos -= self.getValueAtX(xPos)
        
        request.session[request.POST['xID']] = xPos
        request.session[request.POST['yID']] = yPos
        
        # Update function at specified index, or else most recent function (functionIndex is -1)
        self.functions[functionIndex].setFunctionParamsFromRequest(request)
        
        return request
    
    def getFunctionsParamsAsArray(self):
        functionsParamsArray = []
        functionsParamsArraySlices = []
        for function in self.functions:
            functionParamsArray = function.getFunctionParamsAsArray()
            functionsParamsArray.extend(functionParamsArray)
            functionsParamsArraySlices.append(len(functionParamsArray))
        return (functionsParamsArray, functionsParamsArraySlices)
    
    def setFunctionsParamsFromArray(self, params, slices):
        pointer = 0
        counter = 0
        for counter in range(len(slices)):
            functionParamsArray = params[pointer : pointer + slices[counter]]
            self.functions[counter].setFunctionParamsFromArray(functionParamsArray)
            pointer += slices[counter]
            counter += 1
        
    def getFunctionsParamsFromArray(self, params, slices):
        functionsParams = []
        
        pointer = 0
        counter = 0
        for counter in range(len(slices)):
            functionParamsArray = params[pointer : pointer + slices[counter]]
            functionsParams.append(self.functions[counter].getFunctionParamsFromArray(functionParamsArray))
            pointer += slices[counter]
            counter += 1
        
        return functionsParams


    def getFitFunctionInfos(self, mpfitResult):
        params, slices = self.getFunctionsParamsAsArray()
        fitFunctionInfos = []
        pointer = 0
        counter = 0
        for function in self.functions:
            mpfitFunctionResult = dict(params=mpfitResult.params[pointer : pointer + slices[counter]],
                                       perror=mpfitResult.perror[pointer : pointer + slices[counter]])
            
            # Map sigfigs
            fitFunctionParams       = function.getFunctionParamsFromArray([sigfig(x, 6) for x in mpfitFunctionResult['params']])
            fitFunctionParamsErr    = function.getFunctionParamsFromArray([sigfig(x, 2) for x in mpfitFunctionResult['perror']])
            fitFunctionParamsArray  = paramsJoin(fitFunctionParams, fitFunctionParamsErr)       
            
            fitFunctionInfo = { 'fitFunctionParams': fitFunctionParams, 'fitFunctionParamsErr': fitFunctionParamsErr,
                                'fitFunctionParamsArray': fitFunctionParamsArray }
            fitFunctionInfo.update(function.getJSON())
            
            fitFunctionInfos.append(fitFunctionInfo)
            
            pointer += slices[counter]
            counter += 1
        
        print fitFunctionInfos
        print '*****'
        print
        
        return fitFunctionInfos
    

class Function(object):
    def __init__(self, functionID=0):
        self.functionID = functionID
        self.functionName = ''
        self.fitInstructions = deque()
        self.functionParams = FunctionParamGroup()
    
    def __repr__(self):
        return str(self.getJSON())
        
    def function(self, Domain, Range):
        return None
    
    def getJSON(self):
        return { 'functionID': self.functionID, 'functionName': self.functionName, 'functionParams': self.functionParams.asArray() }
    
    def setFunctionParamsFromRequest(self, request):
        for (paramName, value) in self.functionParams.params.items():
            if request.session.has_key(paramName):
                self.functionParams.set(paramName, request.session[paramName])
                
    def setFunctionParamsFromDict(self, functionParamsDict):
        for (paramName, value) in self.functionParams.params.items():
            if functionParamsDict.has_key(paramName):
                self.functionParams.set(paramName, functionParamsDict[paramName])
    
    def setFunctionParamsFromArrayOfDicts(self, functionParamsArrayOfDicts):
        for functionParam in functionParamsArrayOfDicts:
            if self.functionParams.params.has_key(functionParam['paramName']):
                self.functionParams.set(functionParam['paramName'], functionParam['value'])
    
    
    def setFunctionParamsFromArray(self, functionParamsArray):
        functionParamsDict = self.getFunctionParamsFromArray(functionParamsArray)
        self.setFunctionParamsFromDict(functionParamsDict)
    
    def getFunctionParamsAsArray(self):
        pass
    
    def getFunctionParamsFromArray(self, params):
        pass
        

    def createFunction(self, xData, yData):
        functionDomain    = N.arange(min(xData), max(xData), abs(max(xData) - min(xData)) / 200.)
        functionDataRange = N.arange(min(yData), max(yData), abs(max(yData) - min(yData)) / 200.)
        
        functionRange = self.function(functionDomain, functionDataRange)
        functionData = zip_(functionDomain, functionRange)
        
        functionY = self.function(xData, yData)
        functionResiduals = N.subtract(functionY, yData)
        functionResidualData = zip_(xData, functionResiduals)

        print
        print functionRange
        print functionY
        print
        print '--------'
        print
        
        functionObj = dict(fit=functionData, resid=functionResidualData,
                           functionID=self.functionID, functionName=self.functionName, functionParams=self.functionParams)
        return functionObj
          
    def chisq(self, xData, yData, yErr):
        yCalc = self.function(xData, yData)
        
        yErr_temp = copy.deepcopy(yErr)
    #    zero_loc = N.where(yErr == 0)[0]
    #    if len(zero_loc) != 0:
    #        yErr_temp[zero_loc] = 1.0
        chi = ((yData - yCalc) / yErr_temp) ** 2
        
        return chi.sum() / (len(yData) - len(self.functionParams))
        


# LINEAR

class Linear(Function):
    def __init__(self):
        super(Linear, self).__init__()
        self.functionID = 1
        self.functionName = 'Linear'
        self.fitInstructions  = deque([
                                    { 'dataType': 'askPoint', 'xID': 'X1', 'yID': 'Y1',
                                      'messageTitle': 'Step 1', 'messageText': 'Please click on the first point' },
                                    { 'dataType': 'askPoint', 'xID': 'X2', 'yID': 'Y2',
                                      'messageTitle': 'Step 2', 'messageText': 'Please click on the second point' }
                                ])
        self.functionParams = FunctionParamGroup([ 'slope', 'yInter' ])

    def function(self, Domain, Range):
        (slope, yInter) = (self.functionParams.get('slope'), self.functionParams.get('yInter'))
        
        if slope is None or yInter is None:
            pass
        else:
            return N.multiply(slope, Domain) + yInter
    
    def getFunctionParamsAsArray(self):
        return [ self.functionParams.get('slope'), self.functionParams.get('yInter') ]
    
    def getFunctionParamsFromArray(self, params):
        return { 'slope': params[0], 'yInter': params[1] }

    def setFunctionParamsFromRequest(self, request):
        super(Linear, self).setFunctionParamsFromRequest(request)

        if request.session.has_key('Y2'):
            self.setFunctionParamsFromTwoPoints(request.session)

    def setFunctionParamsFromTwoPoints(self, twoPoints):
        slope = N.divide(twoPoints['Y2'] - twoPoints['Y1'], twoPoints['X2'] - twoPoints['X1'])
        yInter = twoPoints['Y1'] - twoPoints['X1'] * slope
        self.functionParams.set('slope', slope)
        self.functionParams.set('yInter', yInter)

class LinearDrag(Linear):
    def __init__(self):
        super(LinearDrag, self).__init__()
        self.functionID = 2
        self.functionName = 'Linear drag'
        self.fitInstructions  = deque([
                                    { 'dataType': 'askDrag', 'xIDstart': 'X1', 'yIDstart': 'Y1', 'xIDend': 'X2', 'yIDend': 'Y2',
                                      'messageTitle': 'Step 1', 'messageText': 'Please drag from the first point to the second point' }
                                ])


# GAUSSIAN

class Gaussian(Function):
    def __init__(self):
        super(Gaussian, self).__init__()
        self.functionID = 11
        self.functionName = 'Gaussian'
        self.fitInstructions  = deque([
                                    { 'dataType': 'askPoint', 'xID': 'peakX', 'yID': 'peakY',
                                      'messageTitle': 'Step 1', 'messageText': 'Please click on the peak of the data' },
                                    { 'dataType': 'askPoint', 'xID': 'widthX', 'yID': 'widthY',
                                      'messageTitle': 'Step 2', 'messageText': 'Please click on the width of the data' }
                                ])
        self.functionParams = FunctionParamGroup([ 'peakX', 'peakY', 'FWHM' ])

    def function(self, Domain, Range):
        (peakX, peakY, FWHM) = (self.functionParams.get('peakX'), self.functionParams.get('peakY'), self.functionParams.get('FWHM'))
        bkgdY = 0
        
        if peakX is None or FWHM is None:
            pass
        else:
            stdDev = FWHM / 2 / N.sqrt(2 * N.log(2))
            return bkgdY + (peakY - bkgdY) * N.exp(- N.power(N.subtract(Domain, peakX), 2) / 2 / N.power(stdDev, 2))
    
    def getFunctionParamsAsArray(self):
        return [ self.functionParams.get('peakX'), self.functionParams.get('peakY'),
                 self.functionParams.get('FWHM') ]
        
    def getFunctionParamsFromArray(self, params):
        return { 'peakX': params[0], 'peakY': params[1], 'FWHM': params[2] }
    
    def setFunctionParamsFromRequest(self, request):
        super(Gaussian, self).setFunctionParamsFromRequest(request)
        
        if request.session.has_key('widthX'):
            FWHM = 2 * N.abs(request.session['widthX'] - request.session['peakX'])
        self.functionParams.set('FWHM', FWHM)
        
class GaussianDrag(Gaussian):
    def __init__(self):
        super(GaussianDrag, self).__init__()
        self.functionID = 12
        self.functionName = 'Gaussian drag'
        self.fitInstructions  = deque([
                                    { 'dataType': 'askPoint', 'xID': 'peakX', 'yID': 'peakY',
                                      'messageTitle': 'Step 1', 'messageText': 'Please click on the peak of the data' },
                                    { 'dataType': 'askDrag', 'xIDstart': 'widthYst', 'yIDstart': 'widthYst', 'xIDend': 'widthX', 'yIDend': 'widthY',
                                      'messageTitle': 'Step 2', 'messageText': 'Please drag on the width of the data' }
                                ])


# LORENTZIAN

class Lorentzian(Function):
    def __init__(self):
        super(Lorentzian, self).__init__()
        self.functionID = 21
        self.functionName = 'Lorentzian'
        self.fitInstructions  = deque([
                                    { 'dataType': 'askPoint', 'xID': 'peakX', 'yID': 'peakY',
                                      'messageTitle': 'Step 1', 'messageText': 'Please click on the peak of the data' },
                                    { 'dataType': 'askPoint', 'xID': 'widthX', 'yID': 'widthY',
                                      'messageTitle': 'Step 2', 'messageText': 'Please click on the width of the data' }
                                ])
        self.functionParams = FunctionParamGroup([ 'peakX', 'peakY', 'FWHM' ])

    def function(self, Domain, Range):
        (peakX, peakY, FWHM) = (self.functionParams.get('peakX'), self.functionParams.get('peakY'), self.functionParams.get('FWHM'))
        bkgdY = 0
        
        if peakX is None or FWHM is None:
            pass
        else:
            gamma = FWHM
            return bkgdY + (peakY - bkgdY) * N.divide(N.power(gamma, 2), N.power(N.subtract(Domain, peakX), 2) + N.power(gamma, 2))

    def getFunctionParamsAsArray(self):
        return [ self.functionParams.get('peakX'), self.functionParams.get('peakY'), self.functionParams.get('FWHM') ]
        
    def getFunctionParamsFromArray(self, params):
        return { 'peakX': params[0], 'peakY': params[1], 'FWHM': params[2] }
        
    def setFunctionParamsFromRequest(self, request):
        super(Lorentzian, self).setFunctionParamsFromRequest(request)
        
        if request.session.has_key('widthX'):
            FWHM = N.abs(request.session['widthX'] - request.session['peakX'])
        self.functionParams.set('FWHM', FWHM)

class LorentzianDrag(Lorentzian):
    def __init__(self):
        super(LorentzianDrag, self).__init__()
        self.functionID = 22
        self.functionName = 'Lorentzian drag'
        self.fitInstructions  = deque([
                                    { 'dataType': 'askPoint', 'xID': 'peakX', 'yID': 'peakY',
                                      'messageTitle': 'Step 1', 'messageText': 'Please click on the peak of the data' },
                                    { 'dataType': 'askDrag', 'xIDstart': 'widthYst', 'yIDstart': 'widthYst', 'xIDend': 'widthX', 'yIDend': 'widthY',
                                      'messageTitle': 'Step 2', 'messageText': 'Please drag on the width of the data' }
                                ])


def zip_(l1, l2):
    return [list(elem) for elem in zip(l1, l2)]
    

def paramsJoin(d1, d2):
    """Joins each parameter with its error and name as a dict"""
    n = []
    for (key, value) in d1.items():
        n.append({ 'name': key, 'value': value, 'err': d2[key] })
    return n



def sigfig(x, n=6):
     if n < 1:
         raise ValueError("number of significant digits must be >= 1")
     return "%.*e" % (n - 1, x)