#  LMFit.py
#  M.Lampton UCB SSL 2015
#  Demonstrate fitting a 2D Gaussian PSF to a bunch of data.
#  Application will be to getting star centroids from robot fiber dance.
#
#  Method: Levenberg Marquardt nonlinear least squares fitter for parms
#  PSF parms are: Xcenter, Ycenter, FWHM, Signal, Background/fiber.
#  Units for x, y, r: microns
#  Units for signal: total electrons from star in exposure
#  Units for background: electrons per exposure per fiber.
#  Fit is linear in signal and bkg, nonlinear in x, y, r. 
#
#  Also evaluates the uncertainties on these five best fit parms.

import math
import numpy as np
from scipy import linalg


#  Declare all constants and constant arrays.
#  These appear only on RHSides and are therefore automatically global.
#  I use all upper case for these global constants.
 
RFIBER = 55.0;      # fiber radius microns
READNOISE = 100.0   # rms electrons per fiber
#  Dark current is included within background parameter.

#  List eachfiber center {x,y} coordinates in microns
#  This is a tidy 5x5 array but could be ratty and irregular
PIXXY = np.array(
        [[-200.0,  -200.0],
         [-100.0,  -200.0],
         [   0.0,  -200.0],
         [ 100.0,  -200.0],
         [ 200.0,  -200.0],
         [-200.0,  -100.0],
         [-100.0,  -100.0],
         [   0.0,  -100.0],
         [ 100.0,  -100.0],
         [ 200.0,  -100.0],
         [-200.0,    -0.0],
         [-100.0,    -0.0],
         [   0.0,    -0.0],
         [ 100.0,    -0.0],
         [ 200.0,    -0.0],
         [-200.0,  +100.0],
         [-100.0,  +100.0],
         [   0.0,  +100.0],
         [ 100.0,  +100.0],
         [ 200.0,  +100.0],
         [-200.0,  +200.0],
         [-100.0,  +200.0],
         [   0.0,  +200.0],
         [ 100.0,  +200.0],
         [ 200.0,  +200.0]])
NPTS = len(PIXXY)
print "Number of data {xy} positions is...", NPTS

#  List the averaging points with one fiber  
#  I chose 19 points within unit radius circle
AVEXY = np.array(
        [[ 0.000000,  0.000000],
         [ 0.346410,  0.200000],
         [ 0.000000,  0.400000],
         [-0.346410,  0.200000],
         [-0.346410, -0.200000],
         [-0.000000, -0.400000],
         [ 0.346410, -0.200000],
         [ 0.800000,  0.000000],
         [ 0.692820,  0.400000],
         [ 0.400000,  0.692820],
         [ 0.000000,  0.800000],
         [-0.400000,  0.692820],
         [-0.692820,  0.400000],
         [-0.800000,  0.000000],
         [-0.692820, -0.400000],
         [-0.400000, -0.692820],
         [-0.000000, -0.800000],
         [ 0.400000, -0.692820],
         [ 0.692820, -0.400000]])
NAVE = len(AVEXY)
print "Number of averaging points within a fiber is...", NAVE

#  True parms........   Xc,um  Yc,um  FWHM,um  Signal   Bkg/fiber
#  Linear dimensions are in microns at the focal surface.
#  Signal is total electrons from target star, entire focal plane.
#  Background is electrons per fiber: night sky and dark current.
TRUEP     = np.array([  10.0,  10.0,  120.0,   12000.0,  100.0])
NPARMS    = len(TRUEP)


#  Set up the adjustable parameters and arrays.
#  These appear on both LHSides and RHSides; 
#  so they must be declared global by each function using them.
#
#  Start fit at....     Xc,um  Yc,um  FWHM,um  Signal   Bkg/fiber
gloparms  = np.array([   0.0,   0.0,   90.0,   25000.0,    0.0])
glodata   = np.zeros(NPTS)           # fake constant data
glonoise  = np.zeros(NPTS)           # estimated rms noise per fiber
gloresid  = np.zeros(NPTS)           # weighted discrepancy of fit
glojacob  = np.zeros((NPTS,NPARMS))  # rows,cols of jacobian matrix
glosigma  = np.zeros(NPARMS)         # derived parameter rms errors



#  Declare all needed functions.

def psf(r, fwhm):
# Here I model our PSF as a 2D Gaussian with given FWHM.
# This is the density of electrons per square microns, per event.
    s = fwhm/2.35
    if s <= 0:
      return 0.
    arg = 0.5*(r/s)**2
    if arg > 15:
      return 0.
    return math.exp(-arg)/(2.*math.pi*s**2)
    


def putNoiselessData(parms, result):
# Here I integrate the PSF falling into each fiber.
    dArea = (math.pi * RFIBER**2)/NAVE;   # square microns
    for i in range(0, NPTS):
        sum = parms[4]
        for j in range(0, NAVE):
            dx = PIXXY[i][0] - parms[0] + RFIBER*AVEXY[j][0]
            dy = PIXXY[i][1] - parms[1] + RFIBER*AVEXY[j][1]
            r = math.sqrt(dx**2 + dy**2)
            sum += parms[3] * dArea * psf(r, parms[2])
        result[i] = sum
    return
    
    

def putNoisyData():
# Here I add Poisson-rule Gaussian noise to each datum.
# Also fills in the glonoise vector. 
    global glodata
    global glonoise
    putNoiselessData(TRUEP, glodata) 
    for i in range(0, NPTS):
        read = np.random.normal() * READNOISE
        shot = np.random.normal() * math.sqrt(glodata[i])
        glodata[i] += read + shot; 
        glonoise[i] = math.sqrt(READNOISE**2 + glodata[i])
    return
    


def chisq(parms):
# Here I evaluate chi squared for any given parameter vector.
# Incidentally also compute the residuals vector, for the jacobian.
# This is based on precalculated data[] and rtwt[] vectors.
    global glodata
    global glortwt
    global glonoise
    model = np.zeros(NPTS)
    putNoiselessData(parms, model)
    sumsq = 0.0
    for i in range(0, NPTS):
        gloresid[i] = (model[i] - glodata[i])/glonoise[i]
        sumsq += gloresid[i]**2
    return sumsq
    


def func(which, bestfit, dx):
# Here I create a 1D explorer of the multidimensional chisq function.
# "Which" says which parameter to deviate by "dx" from bestfit parms[].
# I subtract 1 to make easy rootfinding for chisq rise = 1. 
    trial = np.zeros(NPARMS)
    for i in range(0, NPARMS):
        trial[i] = bestfit[i]
    trial[which] += dx
    return chisq(trial) - chisq(bestfit) - 1
    
    
def getSigmas(bestfit, sig):
# Here I evaluate the parameter sigmas using curvature along each axis.
    delta = 1.0
    for i in range(0, NPARMS):
        fa = func(i, bestfit, -delta)
        fb = func(i, bestfit, 0.)
        fc = func(i, bestfit, delta)
        denom = math.sqrt(0.5*fa + 0.5*fc - fb)
        sig[i] = delta/denom
    return
    
    
#-----Interface service functions for LM minimization-----------
#----Interface isolation is important with separate classes------
#------Unimportant in Python with everything in one class---


def dNudge(dp):
# This one nudges the current global parms[] by a given dp,
# and reports the resulting chisq.
    global gloparms
    # optional printout of the dp vector...
    # np.set_printoptions(precision=4,suppress=True)  # suppress e-notation
    # print "showing dNudge() dp vector...."
    # print dp
    for i in range(0, NPARMS):
        gloparms[i] += dp[i]
    return chisq(gloparms)
    


def dGetResid(i):
# This one allows LM to gain access to the current gloresid[] vector.
# Index "i" is in range 0, NPTS
    global gloresid
    return gloresid[i]
    
    

def dGetJacob(i, j):
# Access to Jacobian matrix of partial derivatives
# index i<NPTS; j<NPARMS
    global glojacob
    return glojacob[i][j]
   
    

def buildJacobian():
# This one permits LM to request an entirely new global jacobian.
# Method: two sided finite difference. 
    DELTAP = 1.0
    FACTOR = 0.5/DELTAP            # converts two sided difference into derivative
    for j in range(0, NPARMS):
        deltap = np.zeros(NPARMS)  # zero the deltap
        deltap[j] = DELTAP         # choose "j" positive direction
        dNudge(deltap)             # new parms; get resid at pplus
        for i in range(0, NPTS):
            glojacob[i][j] = dGetResid(i)
            
        deltap[j] = -2*DELTAP      # choose "j" negative direction
        dNudge(deltap)             # new parms; get resid at pminus
        for i in range(0, NPTS):
            glojacob[i][j] -= dGetResid(i)
        deltap[j] = DELTAP         # done with this axis
        dNudge(deltap)             # so move parms back home
        for i in range(0, NPTS):   # convert finite diffs to derivs
            glojacob[i][j] *= FACTOR
    # optional printout of the jacobian matrix...
    # np.set_printoptions(precision=4,suppress=True)  # suppress e-notation
    # print glojacob
    return
    


#-----Define the LM host and iteration methods-----------
#-----These have their own working constants and arrays-----
#----These appear only on RHSides and are therefore assumed global------
#----In contrast, items that appear on LHSides are assumed local-------

LMITER      = 100
LMBOOST     = 2.0
LMSHRINK    = 0.1
LAMZERO     = 0.01
LAMMAX      = 1E3
LMTOL       = 1E-6

#---LM variables must be read/write hence declared global------------


glolam      = LAMZERO   # users must declare this global

def bLMiter():          
# Does one LM iteration; uses scipy linalg.inv(A) matrix inverter.
# This is called by LM() which is the host of the process. 
    global glolam
    delta   = np.zeros(NPARMS)
    beta    = np.zeros(NPARMS)
    alpha   = np.zeros((NPARMS,NPARMS)) 
    amatrix = np.zeros((NPARMS,NPARMS))
    delta   = np.zeros(NPARMS)   # no changes yet
    sos = dNudge(delta);         # get sos at beginning of this iteration
    print "bLMiter starting with sos = %.2f"  % sos
    sosprev = sos
    
    buildJacobian()              # start with a fresh jacobian.

    for k in range(0, NPARMS):   # get the downhill gradient
        beta[k] = 0.
        for i in range(0, NPTS):
            beta[k] -= dGetResid(i)* dGetJacob(i,k)
            
    # print "Printing beta...."
    # np.set_printoptions(precision=2,suppress=True)  # suppress e-notation
    # print beta

    for k in range(0, NPARMS):   # get the curvature matrix alpha
        for j in range(0, NPARMS):
            alpha[j][k] = 0.
            for i in range(0, NPTS):
                alpha[j][k] += dGetJacob(i,j)*dGetJacob(i,k)
    while True:                              # get a downhill step
        for k in range(0, NPARMS):
            for j in range(0, NPARMS):
                amatrix[j][k] = alpha[j][k]  # copy alpha
                if j == k:
                    amatrix += glolam        # apply some damping
                    
        amatrix = linalg.inv(amatrix)        # invert
        
        for k in range(0, NPARMS):
            delta[k] = 0.
            for j in range(0, NPARMS):
                delta[k] += amatrix[j][k]*beta[j]
                
        sos = dNudge(delta)               # try it out
        # print "trying sos = %.2f" % sos
        rrise = (sos - sosprev)/(1+sos)
        if rrise <= 0:                       # got a downhill step!
            glolam *= LMSHRINK               # shrink damping
            # print "nice drop, shrinking"
            break                            # exit this iteration
        
        for q in range(0, NPARMS):
            delta[q] *= -1.                  # reverse course
            dNudge(delta)
        if rrise < LMTOL:                    # finished
            # print "tiny rise; exitting"
            break                            # exit this iteration
            
        glolam *= LMBOOST                    # jack up the damping
        if glolam > LAMMAX:                  # ridiculous? bail out
            # print "exceeded LAMMAX; exitting"
            break
        print "trying a larger lambda"
    done = (rrise>-LMTOL) or (glolam>LAMMAX)
    return done
    

def LM():   
# Here is the host for the LM iterations.
    global glolam
    glolam = 0.001
    niter = 0
    while True:        # forever
        niter += 1
        if bLMiter():   # true when done
           break
    print "Done with iterations:  niter=%d " % niter
    return



print "\nStarting LMfit"

print "\nPrinting TRUEP[]..."
print "    X,um    Y,um   FWHMum  Sig,e   Bkg,e"
np.set_printoptions(precision=1,suppress=True)  # suppress e-notation
print TRUEP

print "\nPrinting initial guess parms[]..."
print "    X,um    Y,um   FWHMum  Sig,e   Bkg,e"
np.set_printoptions(precision=1,suppress=True)  # suppress e-notation
print gloparms


print "\nCreate and show noisy data[] in 5x5 format..."
print "Idea is to show what a robotically-stepped fiber image looks like."
print "Expect to get roughly half of the total starlight into central fiber,"
print "declining as we examine fiber locations further from the center."
putNoisyData()                    # puts glodata[] and glonoise 
temp = glodata.reshape((5,5))     # reshape to make a 5x5 image
np.set_printoptions(precision=1)  # drops trailing zeros, uses tabs to set columns
print temp

LM()

getSigmas(gloparms, glosigma)

print "\nListing best fit parms and their errors"
print '   Xcenter {:8.2f}  {:8.2f}'.format(gloparms[0], glosigma[0])
print '   Ycenter {:8.2f}  {:8.2f}'.format(gloparms[1], glosigma[1])
print '      FWHM {:8.2f}  {:8.2f}'.format(gloparms[2], glosigma[2])
print '    Signal {:8.2f}  {:8.2f}'.format(gloparms[3], glosigma[3])
print '       Bkg {:8.2f}  {:8.2f}'.format(gloparms[4], glosigma[4])