ABACUS/src/FITTING/mrq.cc

/**********************************************************

This software is part of J.-S. Caux's ABACUS library.

Copyright (c) J.-S. Caux.

-----------------------------------------------------------

File:  mrq.cc

Purpose:  mrqmin and mrqcof algorithms

***********************************************************/

#include "ABACUS.h"

using namespace std;

namespace ABACUS {

  void mrqmin (Vect_DP& x, Vect_DP& y, Vect_DP& sig, Vect_DP& a,
	       Vect<bool>& ia, SQMat_DP& covar, SQMat_DP& alpha, DP& chisq,
	       void funcs(const DP, Vect_DP&, DP&, Vect_DP&), DP& alambda)
  {
    // Levenberg-Marquardt method.  See NRC++ p.691.

    static int mfit;
    static DP ochisq;
    int j, k, l;

    int ma = a.size();
    static SQMat_DP oneda (ma);
    static Vect_DP atry(ma);
    static Vect_DP beta(ma);
    static Vect_DP da(ma);

    if (alambda < 0.0) {
      mfit = 0;
      for (j = 0; j < ma; j++)
	if (ia[j]) mfit++;
      alambda = 0.001;
      mrqcof (x, y, sig, a, ia, alpha, beta, chisq, funcs);
      ochisq = chisq;
      for (j = 0; j < ma; j++) atry[j] = a[j];
    }

    SQMat_DP temp (mfit);

    for (j = 0; j < mfit; j++) {
      for (k = 0; k < mfit; k++) covar[j][k] = alpha[j][k];
      covar[j][j] = alpha[j][j] * (1.0 + alambda);
      for (k = 0; k < mfit; k++) temp[j][k] = covar[j][k];
      oneda[j][0] = beta[j];
    }

    gaussj (temp, oneda);
    for (j = 0; j < mfit; j++) {
      for (k = 0; k < mfit; k++) covar[j][k] = temp[j][k];
      da[j] = oneda[j][0];
    }

    if (alambda == 0.0) {
      covsrt (covar, ia, mfit);
      covsrt (alpha, ia, mfit);
      return;
    }

    for (j = 0, l = 0; l < ma; l++)
      if (ia[l]) atry[l] = a[l] + da[j++];

    mrqcof (x, y, sig, atry, ia, covar, da, chisq, funcs);

    if (chisq < ochisq) {
      alambda *= 0.1;
      ochisq = chisq;
      for (j = 0; j < mfit; j++) {
	for (k = 0; k < mfit; k++) alpha[j][k] = covar[j][k];
	beta[j] = da[j];
      }
      for (l = 0; l < ma; l++) a[l] = atry[l];
    } else {
      alambda *= 10.0;
      chisq = ochisq;
    }
  }

  void mrqcof (Vect_DP& x, Vect_DP& y, Vect_DP& sig, Vect_DP& a,
	       Vect<bool>& ia, SQMat_DP& alpha, Vect_DP& beta, DP& chisq,
	       void funcs (const DP, Vect_DP&, DP&, Vect_DP&))
  {
    int i, j, k, l, m, mfit = 0;
    DP ymod, wt, sig2i, dy;

    int ndata = x.size();
    int ma = a.size();

    Vect_DP dyda (ma);

    for (j = 0; j < ma; j++)
      if (ia[j]) mfit++;
    for (j = 0; j < mfit; j++) {
      for (k = 0; k <= j; k++) alpha[j][k] = 0.0;
      beta[j] = 0.0;
    }

    chisq = 0.0;
    for (i = 0; i < ndata; i++) {
      funcs (x[i], a, ymod, dyda);
      sig2i = 1.0/(sig[i] * sig[i]);
      dy = y[i] - ymod;
      for (j = 0, l = 0; l < ma; l++) {
	if (ia[l]) {
	  wt = dyda[l] * sig2i;
	  for (k = 0, m = 0; m < l+1; m++)
	    if (ia[m]) alpha[j][k++] += wt * dyda[m];
	  beta[j++] += dy * wt;
	}
      }
      chisq += dy * dy * sig2i;
    }
    for (j = 1; j < mfit; j++)
      for (k = 0; k < j; k++) alpha[k][j] = alpha[j][k];
  }

}