/** LMdemo Levenberg Marquardt demonstrator
  *
  * This one file contains four classes:
  *     LMdemo        contains just main()
  *     RosenHost     defines Rosenbrock demo problem.
  *     LMhost        interfaces methods needed by LM.
  *     LM            is the Levenberg-Marquardt routine.
  *
  *  M.Lampton UCB SSL (c) 1996; Java edition 2007
  */
public class LMdemo
{
    public static void main(String args[])
    {
        RosenHost myHost = new RosenHost(); 
    }
}




//---------------------RosenHost-----------------------

class RosenHost implements LMhost
// Demonstrates Rosenbrock's parabolic valley
// Triggers LM by instantiating an object of class LM
// Supplies five callback methods mandated by interface LMhost
{
    //--------constants---------------
    protected final double SIGMA  = 1e-6;
    protected final double DELTAP = 1e-6;
    protected final double BIGVAL = 9.876543E+210; 
    protected int NPTS, NPARMS; 

    //--------fields------------------
    private double resid[]; 
    private double jac[][];  
    private double parms[];

    public RosenHost()
    {
        //-----set up constants----
        NPTS = 2; 
        NPARMS = 2; 

        //----set up working arrays---------
        resid = new double[NPTS]; 
        jac = new double[NPTS][NPARMS]; 
        parms = new double[NPARMS]; 
        parms[0] = -1.2;  // Rosenbrock starting point
        parms[1] = +1.0;  // Rosenbrock starting point
        for (int i=0; i<NPARMS; i++)
          System.out.println("Start parm["+i+"] = "+parms[i]); 

        //----run the minimization-------------
        System.out.println("TMAhost is starting LM");
        LM myLM = new LM(this, NPARMS, NPTS); 
        System.out.println("TMAhost is finished with LM");  
        for (int i=0; i<NPARMS; i++)
          System.out.println("End parm["+i+"]   = "+parms[i]); 
    }

    public double dComputeResid()
    // This is Rosenbrock's banana.
    // Evaluates residual vector for parms[].
    // Returns sum-of-squares. 
    {
        double DIFFICULTY = 100.0; 
        resid[0] = DIFFICULTY * (parms[1] - parms[0]*parms[0]); 
        resid[1] = 1.0 - parms[0]; 
        return resid[0]*resid[0] + resid[1]*resid[1]; 
    }

    public double dNudge(double dp[])
    // Allows LM to modify 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 dComputeResid(); 
    }

    public boolean bBuildJacobian()
    // Allows LM to compute a new Jacobian.
    // Uses current parms[] and two-sided finite difference.
    // If current parms[] is bad, returns false.  
    {
        double delta[] = new double[NPARMS];
        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); // resid at pplus
            if (d==BIGVAL)
            {
                System.out.println("Bad dBuildJacobian() exit 2"); 
                return false;  
            }
            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); // resid at pminus
            if (d==BIGVAL)
            {
                System.out.println("Bad dBuildJacobian() exit 3"); 
                return false;  
            }

            for (int i=0; i<NPTS; i++)
              jac[i][j] -= dGetResid(i);  // fetches resid[]

            for (int i=0; i<NPTS; i++)
              jac[i][j] *= FACTOR;

            for (int k=0; k<NPARMS; k++)
              delta[k] = (k==j) ? DELTAP : 0.0;

            d = dNudge(delta);  
            if (d==BIGVAL)
            {
                System.out.println("Bad dBuildJacobian() exit 4"); 
                return false;  
            }
        }
        return true; 
    }

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

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





//---LMhost interface class declares five abstract methods---------

interface LMhost
{
    double dComputeResid();
    // Allows LM to evaluate a starting point. 

    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 here. 

    boolean bBuildJacobian();
    // Allows LM to request a new Jacobian.

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

    double dGetJac(int i, int j);
    // Allows LM to access 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 five host methods:
  *    double dComputeResid();    Returns sos, or BIGVAL if parms failed.
  *    double dNudge(dp);         Moves parms, builds resid[], returns sos.
  *    boolean bBuildJacobian();  Builds Jacobian, returns false if parms NG.
  *    double dGetJac(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 LAMBDAZERO = 0.001;    // initial damping
    private final double LAMBDAMAX  =  1E9;     // max damping
    private final double LMTOL      = 1E-12;    // exit tolerance
    private final double BIGVAL = 9.876543E+210; 
 
    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
    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 = LAMBDAZERO; 
        int niter = 0; 
        boolean done = false; 
        do
        {
            done = bLMiter();
            niter++;
        } 
        while (!done && (niter<LMITER));
    }

    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.
    {
        sos = myH.dComputeResid();
        if (sos==BIGVAL)
        {
           System.out.println("bLMiter finds faulty initial dComputeResid()");
           return false; 
        }
        sosprev = sos;
        System.out.println("bLMiter..sos= "+sos);
        if (!myH.bBuildJacobian())
        {
            System.out.println("bLMiter finds bBuildJacobian()=false"); 
            return false;
        }
        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.dGetJac(i,k);
        }
        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.dGetJac(i,j)*myH.dGetJac(i,k);
          }
        double rrise = 0; 
        do  /// damping loop searches for one downhill step
        {
            // System.out.println("  lambda = "+lambda); 
            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.
            if (sos==BIGVAL)
            {
                System.out.println("LMinner failed SOS step"); 
                return false;            
            }
            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<LAMBDAMAX);
        boolean done = (rrise>-LMTOL) || (lambda>LAMBDAMAX); 
        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--------------------