/** LMfit2D Levenberg Marquardt demonstrator; coordinates are microns
  *
  * Demonstrates fitting a 2D circular Gaussian to noisy data. 
  * Five parameters: Xcenter, Ycenter, FWHM, Signal, Background.
  * X, Y, and FWHM are in microns.
  * Signal is in electrons total for the exposure.
  * Background is electrons per fiber for the exposure.
  * Image is an arbitrary grid of pixels.
  * Data come from a smooth model + random number generator. 
  * Gaussian noise in each image pixel is sqrt(signal+background).
  * Computes best fit parameters using LM.
  * Computes their one-sigma errors by 3-point FD curvature.
  *
  * This one file contains several classes:
  *     LMfitGauss    contains only main(); creates a FitHost.
  *     FitHost       defines the fitting problem.
  *     LMhost        defines interface methods required by LM. 
  *     LM            The Levenberg-Marquardt routine.
  *
  *  M.Lampton UCB SSL (c) 1996; Java edition 2007, 2011, 2015
  *  see M.Lampton, 1997 Computers In Physics v.11 #10 110-115.
  */
public class LMfit2D
{
    public static void main(String args[])
    {
        new FitHost(); 
    }
}





/**  FitHost 
  *  Fits a five parameter Gaussian to pixellized noisy data
  *  With or without subpixel averaging.
  *  Sets up data with noise, and all needed arrays. 
  *  Triggers LM by instantiating an object of class LM
  *  Offers four callback methods mandated by interface LMhost
  */
class FitHost implements LMhost
{
    //--------constants---------------

    protected final double RFIBER    = 55.0;     // microns
    protected final double READNOISE = 100.0;    // rms electrons per fiber
    //----dark current is included within background-----

    //---List the pixel center locations, microns-----
    double pixXY[][] = {{-200.0, -200.0},
                        {-100.0, -200.0},
                        {-000.0, -200.0},
                        {+100.0, -200.0},
                        {+200.0, -200.0},
                        {-200.0, -100.0},
                        {-100.0, -100.0},
                        {-000.0, -100.0},
                        {+100.0, -100.0},
                        {+200.0, -100.0},
                        {-200.0, -000.0},
                        {-100.0, -000.0},
                        {-000.0, -000.0},
                        {+100.0, -000.0},
                        {+200.0, -000.0},
                        {-200.0, +100.0},
                        {-100.0, +100.0},
                        {-000.0, +100.0},
                        {+100.0, +100.0},
                        {+200.0, +100.0},
                        {-200.0, +200.0},
                        {-100.0, +200.0},
                        {-000.0, +200.0},
                        {+100.0, +200.0},
                        {+200.0, +200.0}};
    int NPTS = pixXY.length; 

    //---List the averaging points within one fiber-----------------------
    //---Use a square array for square pixels-----------------------------
    //---Here is a 19 point equal area circular array for round fibers----
    //---These are normalized to a unit radius circular fiber-------------
    double aveXY[][] = {{ 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}};
    int NAVE = aveXY.length;

    
    
    //---Set up true Gaussian PSF a bit off center----------------------
    //---Here, xc  yc and FWHM are in microns-------------------------------
    //---signal is total star electrons in Texposure-----------------
    //---background is from sky and dark current, electrons/fiber--------
    
    protected int NPARMS = 5; 
    //----parms are......       xc,    yc,   FWHM,  signal,  bkg/pix. 
    private double truep[]  = { 10.,   10.,  120.0, 12000.0,  100.0};  // true parms
    private double parms[]  = { 0.0,   0.0,   90.0, 25000.0,    0.0};  // start fit

    private double data[]   = new double[NPTS];          // fake data list
    private double noise[]  = new double[NPTS];          // model rms noise per fiber
    private double resid[]  = new double[NPTS];          // current fit discrepancy
    private double jac[][]  = new double[NPTS][NPARMS];  // jacobian partial derivatives
    private double sigma[]  = new double[NPARMS];        // derived parameter errors

    public FitHost()
    {
        listParms("True Parms", truep); 
        listParms("Start Parms", parms); 
        putNoisyData();  
        listNoisyData(); 
        
        // now invoke the Levenberg-Marquardt fitter....
        LM myLM = new LM(this, NPARMS, NPTS);
        listParms("Best Fit Parms", parms); 
        getSigmas(parms, sigma); 
        listSigma("Parm errors", sigma); 
    }

    //------------mathematical helpers for FitHost--------------

    void listParms(String title, double p[])
    {
        System.out.println(title + "----"); 
        System.out.printf("   X= %.2f %n",p[0]); 
        System.out.printf("   Y= %.2f %n",p[1]); 
        System.out.printf("   F= %.2f %n",p[2]); 
        System.out.printf("   S= %.2f %n",p[3]); 
        System.out.printf("   B= %.2f %n",p[4]); 
    }


    
    private void putNoiselessModel(double p[], double result[])
    // This generates the noiseless fitting function for any model parameters p[].
    // Here I integrate NAVE samples arranged within a circular fiber.
    {
        double dArea = Math.PI*RFIBER*RFIBER/NAVE;  // square microns
        for (int i=0; i<NPTS; i++)
        {
            double sum = p[4];                      // background includes dark
            for (int j=0; j<NAVE; j++)              // add up samples  
            {
                double dx = pixXY[i][0] - p[0] + RFIBER*aveXY[j][0];
                double dy = pixXY[i][1] - p[1] + RFIBER*aveXY[j][1]; 
                double r = Math.sqrt(dx*dx + dy*dy); 
                sum += p[3] * psf(r,p[2]) * dArea; 
            }
            result[i] = sum; 
        }
    }
    
    
    private void putNoisyData()
    // This produces the data array to be fitted and its noise.
    // For actual data it would read in a file of numbers.
    // For this demo I use the noiseless model at truep[],
    // then add shot noise and read noise. 
    {
        putNoiselessModel(truep, data);  // start with true noiseless model
        for (int i=0; i<NPTS; i++)       // add noise
        {
            double read = gaussrand() * READNOISE;
            double shot = gaussrand() * Math.sqrt(data[i]); 
            data[i] += read + shot; 
            noise[i] = Math.sqrt(READNOISE*READNOISE + data[i]); 
        }
    }

    private void listNoisyData()
    {
        System.out.println("Noisy data in a 5x5 map format-----");
        for (int i=0; i<5; i++)
        {
            for (int j=0; j<5; j++)
               System.out.printf("%10.2f", data[i+5*j]); 
            System.out.println();
        }
    }


    private double psf(double r, double fwhm)
    // Given a point at radius "r" and Gaussian 2D "fwhm",
    // both in microns,
    // Returns Gaussian 2D density normalized to unit total integral. 
    // This is electrons per square micron on the focal plane.
    // Its coefficient should be total electrons arriving. 
    // The yield will integrate this function over dArea of fiber. 
    {
        double s = fwhm/2.35; 
        if (s <= 0.0)
          return 0.0;
        double arg = 0.5*r*r/(s*s); 
        if (arg > 15)
          return 0.0; 
        return Math.exp(-arg)/(2*Math.PI*s*s); 
    }

    private double gaussrand()
    // Returns a 1D gaussian variate, unit variance, zero mean
    {
         double sum = -6.0;
         for (int i=0; i<12; i++)
           sum += Math.random(); 
         return sum;
    }

    private double chisq(double p[])
    // Uses any given parameters vector; 
    // Evaluates resid[] vector. 
    // Returns chi squared = weighted-sum-of-squares.
    {
        double model[] = new double[NPTS]; 
        putNoiselessModel(p, model); 
        double sumsq = 0.0; 
        for (int i=0; i<NPTS; i++)
        {
            resid[i] = (model[i] - data[i])/noise[i]; 
            sumsq += resid[i]*resid[i]; 
        }
        return sumsq; 
    }
    
    
    //----methods for getting 1-D parameter errors------------
    
    private double func(int which, double bestfit[], double dx)
    // Evaluates a 1D slice through a chisq() function,
    // Subtracted to reveal rise through chisq()+1 level.
    // This way, func() is ideally set up for rootfinders. 
    {
        double p[] = new double[NPARMS]; 
        for (int i=0; i<NPARMS; i++)
          p[i] = bestfit[i]; 
        p[which] += dx; 
        return chisq(p) - chisq(bestfit) - 1; 
    }
        
    private void getSigmas(double bestfit[], double sig[])
    // Gets all sigmas using three point curvature method.
    // curvature = (0.5fa + 0.5fc - fb)/D^2  = A in quadratic
    // but quadratic A = rise/sigma^2 = 1/sigma^2
    // so sigma = Delta / sqrt(0.5fa + 0.5fc - fb)
    {
        double delta = 1.0; 
        for (int which=0; which<NPARMS; which++)
        {
            double fa = func(which, bestfit, -delta); 
            double fb = func(which, bestfit, 0.); 
            double fc = func(which, bestfit, delta); 
            double denom = Math.sqrt(0.5*fa + 0.5*fc - fb); 
            sig[which] = delta / denom; 
        }
    }                
    
    private void listSigma(String title, double s[])
    {
        System.out.println(title + "----"); 
        System.out.printf("  sX= %.2f %n",s[0]); 
        System.out.printf("  sY= %.2f %n",s[1]); 
        System.out.printf("  sF= %.2f %n",s[2]); 
        System.out.printf("  sS= %.2f %n",s[3]); 
        System.out.printf("  sB= %.2f %n",s[4]); 
    }


     

    //------the four mandated interface methods------------

    public double dNudge(double dp[])
    // Allows LM to modify global parms[] and reevaluate.
    // Returns sum-of-squares for nudged params.
    // This is the only place that parms[] are modified.
    // If NADJ<NPARMS, this is the place for your LUT.
    {
        for (int j=0; j<NPARMS; j++)
          parms[j] += dp[j]; 
        return chisq(parms); 
    }

    public void buildJacobian()
    // Allows LM to compute a new Jacobian.
    // Uses current parms[] and two-sided finite difference.
    // DELTAP is the parameter step size for finite differences.
    {
        double delta[] = new double[NPARMS];
        double DELTAP = 1.0;
        double FACTOR = 0.5 / DELTAP; 
        double d=0; 

        for (int j=0; j<NPARMS; j++)
        {
            for (int k=0; k<NPARMS; k++)
              delta[k] = (k==j) ? DELTAP : 0.0;
            d = dNudge(delta);        // parms to pplus, new resid.
            for (int i=0; i<NPTS; i++)
              jac[i][j] = dGetResid(i);

            for (int k=0; k<NPARMS; k++)
              delta[k] = (k==j) ? -2*DELTAP : 0.0;
            d = dNudge(delta);        // parms to pminus, new resid.
            for (int i=0; i<NPTS; i++)
              jac[i][j] -= dGetResid(i); 

            for (int k=0; k<NPARMS; k++)
              delta[k] = (k==j) ? DELTAP : 0.0;
            d = dNudge(delta);       // send parms and resid back home. 
            
            for (int i=0; i<NPTS; i++)
              jac[i][j] *= FACTOR;   // convert finite diffs to derivs.
        }
        // Display the jacobian matrix....
        // for (int i=0; i<NPTS; i++)
        // {
        //     for (int k=0; k<NPARMS; k++)
        //       System.out.printf("%12.4f", jac[i][k]);
        //     System.out.println();
        // }
    }

    public double dGetResid(int i)
    // Allows LM to access one element of the resid[] vector. 
    {
        return resid[i];
    }

    public double dGetJacob(int i, int j)
    // Allows LM to access one element of the Jacobian matrix. 
    {
        return jac[i][j]; 
    }
}








/** LMhost interface class declares four abstract methods
  * These allow LM to request numerical work
  * M.Lampton UCB SSL 2005
  */
interface LMhost
{
    double dNudge(double dp[]);
    // Allows myLM.bLMiter to modify parms[] and reevaluate.
    // This is the only modifier of parms[].
    // So, if NADJ<NPARMS, put your LUT into dNudge().

    void buildJacobian();
    // Allows LM to request a new Jacobian.

    double dGetResid(int i);
    // Allows LM to see one element of the resid[] vector. 

    double dGetJacob(int i, int j);
    // Allows LM to see one element of the Jacobian matrix. 
}








/**
  *  class LM   Levenberg Marquardt w/ Lampton improvements
  *  M.Lampton, 1997 Computers In Physics v.11 #10 110-115.
  *
  *  Constructor is used to set up all parms including host for callback.
  *  bLMiter() performs one iteration.
  *  Host arrays parms[], resid[], jac[][] are unknown here.
  *  Callback method uses CallerID to access four host methods:
  *
  *    double dNudge(dp);         Moves parms, builds resid[], returns sos.
  *    boolean buildJacobian();   Builds a new Jacobian.
  *    double dGetJacob(i,j);       Fetches one value of host Jacobian.
  *    double dGetResid(i);       Fetches one value of host residual.
  *
  *  Exit leaves host with optimized parms[]. 
  *
  *  @author: M.Lampton UCB SSL (c) 2005
  */
class LM    
{
    private final int    LMITER     =  100;     // max number of L-M iterations
    private final double LMBOOST    =  2.0;     // damping increase per failed step
    private final double LMSHRINK   = 0.10;     // damping decrease per successful step
    private final double LAMZERO    = 0.01;     // initial damping
    private final double LAMMAX     =  1E3;     // max damping
    private final double LMTOL      = 1E-6;     // exit tolerance
 
    private double sos, sosprev, lambda;

    private LMhost myH = null;    // overwritten by constructor
    private int nadj = 0;         // overwritten by constructor
    private int npts = 0;         // overwritten by constructor

    private double[] delta;       // local parm change
    private double[] beta;        // local
    private double[][] alpha;     // local
    private double[][] amatrix;   // local 

    public LM(LMhost gH, int gnadj, int gnpts)
    // Constructor sets up fields and drives iterations. 
    {
        myH = gH;
        nadj = gnadj;
        npts = gnpts;  
        delta = new double[nadj];
        beta = new double[nadj];
        alpha = new double[nadj][nadj]; 
        amatrix = new double[nadj][nadj];
        lambda = LAMZERO; 
        int niter = 0; 
        boolean done = false; 
        do
        {
            done = bLMiter();
            niter++;
        } 
        while (!done && (niter<LMITER));
        System.out.printf("Niterations = %d     final sos = %.2f \n", niter, sos); 
    }

    private boolean bLMiter( )
    // Each call performs one LM iteration. 
    // Returns true if done with iterations; false=wants more. 
    // Global nadj, npts; needs nadj, myH to be preset. 
    // Ref: M.Lampton, Computers in Physics v.11 pp.110-115 1997.
    {
        for (int k=0; k<nadj; k++)
          delta[k] = 0.0; 
        sos = myH.dNudge(delta);

        sosprev = sos;

        System.out.printf("bLMiter..SumOfSquares= %.2f %n",sos);
        
        myH.buildJacobian();
        
        for (int k=0; k<nadj; k++)      // get downhill gradient beta
        {
            beta[k] = 0.0;
            for (int i=0; i<npts; i++)
              beta[k] -= myH.dGetResid(i)*myH.dGetJacob(i,k);
        }
        
        // System.out.print("Beta = ");
        // for (int k=0; k<nadj; k++)
        //   System.out.printf("   %8.2f", beta[k]);
        // System.out.println(); 
        
        for (int k=0; k<nadj; k++)      // get curvature matrix alpha
          for (int j=0; j<nadj; j++)
          {
              alpha[j][k] = 0.0;
              for (int i=0; i<npts; i++)
                alpha[j][k] += myH.dGetJacob(i,j)*myH.dGetJacob(i,k);
          }
        double rrise = 0; 
        do  /// damping loop searches for one downhill step
        {
            for (int k=0; k<nadj; k++)       // copy and damp it
              for (int j=0; j<nadj; j++)
                amatrix[j][k] = alpha[j][k] + ((j==k) ? lambda : 0.0);

            gaussj(amatrix, nadj);           // invert

            for (int k=0; k<nadj; k++)       // compute delta[]
            {
                delta[k] = 0.0; 
                for (int j=0; j<nadj; j++)
                  delta[k] += amatrix[j][k]*beta[j];
            }
            sos = myH.dNudge(delta);         // try it out.

            rrise = (sos-sosprev)/(1+sos);
            if (rrise <= 0.0)                // good step!
            {
                lambda *= LMSHRINK;          // shrink lambda
                break;                       // leave lmInner.
            }
            for (int q=0; q<nadj; q++)       // reverse course!
              delta[q] *= -1.0;
            myH.dNudge(delta);               // sosprev should still be OK
            if (rrise < LMTOL)               // finished but keep prev parms
              break;                         // leave inner loop
            lambda *= LMBOOST;               // else try more damping.
        } while (lambda<LAMMAX);
        boolean done = (rrise>-LMTOL) || (lambda>LAMMAX); 
        return done; 
    }

    private double gaussj( double[][] a, int N )
    // Inverts the double array a[N][N] by Gauss-Jordan method
    // M.Lampton UCB SSL (c)2003, 2005
    {
        double det = 1.0, big, save;
        int i,j,k,L;
        int[] ik = new int[100];
        int[] jk = new int[100];
        for (k=0; k<N; k++)
        {
            big = 0.0;
            for (i=k; i<N; i++)
              for (j=k; j<N; j++)          // find biggest element
                if (Math.abs(big) <= Math.abs(a[i][j]))
                {
                    big = a[i][j];
                    ik[k] = i;
                    jk[k] = j;
                }
            if (big == 0.0) return 0.0;
            i = ik[k];
            if (i>k)
              for (j=0; j<N; j++)          // exchange rows
              {
                  save = a[k][j];
                  a[k][j] = a[i][j];
                  a[i][j] = -save;
              }
            j = jk[k];
            if (j>k)
              for (i=0; i<N; i++)
              {
                  save = a[i][k];
                  a[i][k] = a[i][j];
                  a[i][j] = -save;
              }
            for (i=0; i<N; i++)            // build the inverse
              if (i != k)
                a[i][k] = -a[i][k]/big;
            for (i=0; i<N; i++)
              for (j=0; j<N; j++)
                if ((i != k) && (j != k))
                  a[i][j] += a[i][k]*a[k][j];
            for (j=0; j<N; j++)
              if (j != k)
                a[k][j] /= big;
            a[k][k] = 1.0/big;
            det *= big;                    // bomb point
        }                                  // end k loop
        for (L=0; L<N; L++)
        {
            k = N-L-1;
            j = ik[k];
            if (j>k)
              for (i=0; i<N; i++)
              {
                  save = a[i][k];
                  a[i][k] = -a[i][j];
                  a[i][j] = save;
              }
            i = jk[k];
            if (i>k)
              for (j=0; j<N; j++)
              {
                  save = a[k][j];
                  a[k][j] = -a[i][j];
                  a[i][j] = save;
              }
        }
        return det;
    }
} //-----------end of class LM--------------------