(*
    Package:        NonlinearFit
    Contains:       Nonlinear model fitting to numerical data
    Method:         Levenberg-Marquardt
    Reference:      Press et al., Numerical Recipes, Cambridge 1986, p. 526
    Written:        by Pekka Janhunen, Finnish Meteorological Institute,
                    Geophysics Dept., email: pjanhune@finsun.csc.fi
    Copyright:      PD software. You are free to use, modify and distribute
                    this software without taking a charge for it.
*)

BeginPackage["NonlinearFit`"]

NonlinearFit::usage =
"NonlinearFit[data,model,params,vars{,options...}] finds values of params
such that the expression model fits to data using the Levenberg - 
Marquardt method. The result is returned as a list of transformation
rules analogous to Solve. Error estimates appear in the result,too.
The meaning of the arguments data and vars is the same as in Fit.
The argument model is any real-valued expression depending on
params and vars. params must be a list of symbols.\n
Supported options: Errors, MaxIterations, InitialGuess, Simplifier,
ErrorNotation, ProgressTrace; help is available on each."

Errors::usage =
"Errors is an option for NonlinearFit that determines the standard deviations
of the data points. Give either a single number or a list of numbers
with length equal to data length. The default is unit standard
deviation for all data points."

MaxIterations::usage =
"MaxIterations is an option for NonlinearFit that determines the maximum
number of iterations allowed. The default value is 20."

InitialGuess::usage =
"InitialGuess is an option for NonlinearFit that determines the initial
values used for the parameters. Give a list of numbers with length
equal to the number of parameters. The default initial guess
is 1.0 for all parameters. Usually you must specify this option
to achieve convergence."

Simplifier::usage =
"Simplifier is an option for NonlinearFit that gives the simplifier function
that is used to simplify the derivatives of the model function before
numerical calculations. Typical choices are Simplify, TrigCanonical
and Together. The default is Identity (no simplification).
A cleaver use of this option can speed up calculations."

ErrorNotation::usage =
"ErrorNotation is an option for NonlinearFit that gives the notation of
the error estimates that appear in the resulting list of transformation
rules. The probability that the true error of each parameter is less
than the given estimate is 68 % if the errors are normally distributed.
The default is Sigma."

ProgressTrace::usage =
"ProgressTrace is an option for NonlinearFit. Set ProgressTrace->True
 to get an informative message on each iteration. Default: False."

Begin["`Private`"]

(*
DiagonalPart[M_List]:=
  Table[If[ii==jj,M[[ii,ii]],0],{ii,1,Length[M]},{jj,1,Length[M]}] /;
  MatrixQ[M] && (Equal @@ Dimensions[M])
*)

DiagonalPart[M_List]:= M IdentityMatrix[Length[M]]

Diagonal[M_List]:= Table[M[[ii,ii]],{ii,1,Length[M]}]/;
  MatrixQ[M] && (Equal @@ Dimensions[M])

EveryQ[expr_,pattern_]:= And @@ (MatchQ[#,pattern]& /@ expr)

MapN[f_,lists__]:= (f @@ #&) /@ Transpose[{lists}]

ReplaceLists[list1_,list2_]:= MapN[#1->#2&,list1,list2]

MarquardtCoefficients \
  [x_List,     (* numerical data abscissas *)
   y_List,     (* numerical data ordinates *)
   stddev_List,(* data standard deviations *)
   params_List,(* numerical parameter values *)
   f_,         (* f[x,params] model function *)
   Df_]:=      (* Df[x,params]={D[f[x,params],par1], D[f[x,params],par2]..}*)
   (* Returns: The list {alpha,beta,chisq} *)
(* Non-iterative version: *)
Block[{dy = y - (N[f[#,params]]&) /@ x,
       dyda = (N[Df[#,params]]&) /@ x,
       invstddev2 = 1/stddev^2},
      { Plus @@ MapN[#1 Outer[Times,#2,#2]&, invstddev2,dyda],
        Plus @@ MapN[#1 #2 #3&, invstddev2,dy,dyda],
        Plus @@ MapN[#1^2 #2&,dy,invstddev2] } ]
(* Iterative version:
Block[
  {alpha=0.,beta=0.,chisq=0.,invstddev2,ymod,dyda,dy,i},
  Do[
     ymod = N[f[x[[i]],params]];
     dyda = N[Df[x[[i]],params]];
     invstddev2 = stddev[[i]]^-2;
     dy = y[[i]] - ymod;
     alpha+= invstddev2 Outer[Times,dyda,dyda];
     beta+= invstddev2 dy dyda;
     chisq+= dy^2 invstddev2,
     {i,Length[x]} ];
  {alpha,beta,chisq}]*)

MarquardtIteration \
  [x_List,      (* numerical data abscissas *)
   y_List,      (* numerical data ordinates *)
   stddev_List, (* data standard deviations *)
   params_List, (* adjustable parameters (what we search for !) *)
   f_,          (* f[x,params] is our model function *)
   Df_,         (* Df[x,params] = {D[f[x,params],par1], D[..,par2],...} *)
   ilambda_]:=  (* the scaling factor *)
    (* Returns: The list {params,chisq,lambda,alpha} *)
    (* Internal variables: beta,oldchisq *)
Block[
  {chisq,dParams,paramsTry,covar,lambda=ilambda},
  If[lambda<0, (* Initialize *)
     lambda=0.001;
     {alpha,beta,chisq} =
       MarquardtCoefficients[x,y,stddev,params,f,Df];
     oldchisq = chisq];
  covar = Inverse[alpha + lambda DiagonalPart[alpha]];
  paramsTry = params + covar . beta;
  {covar,dParams,chisq} =
    MarquardtCoefficients[x,y,stddev,paramsTry,f,Df];
  If[chisq < oldchisq,
     alpha=covar; beta=dParams;  (* Success *)
     {paramsTry,oldchisq=chisq,0.1 lambda,alpha},
     {params,chisq=oldchisq,10.0 lambda,alpha}]  (* Failure *)
]

NonlinearFit::baddata= "`1` is not a real matrix."
NonlinearFit::badvar = "Variable list `1` is not a list of `2` symbols."
NonlinearFit::badpar = "Parameter list `1` is not a list of symbols."

NonlinearFit::baddev = "`1` is not a proper Errors option."
NonlinearFit::short  = "warning: Too short list for Errors option; padded with average."
NonlinearFit::long   = "warning: Too long list for Errors option; truncated."
NonlinearFit::noconv = "No convergence in NonlinearFit."
NonlinearFit::guess  = "warning: `1` is not a proper guess for parameters `2`."

Options[NonlinearFit] =
  {Errors->1., MaxIterations->20, InitialGuess->Automatic,
   Simplifier->Identity, ErrorNotation->Global`Sigma,
   ProgressTrace->False}

nlfit \
  [data_List, model_, paramsymbols_List, vars_List, opts___]:=
Block[
   {x,y,stddev, maxits,guess,params,simp,sigma,tracing,
    X,i,covar,chisq,
    lambda  = -1,
    oldChisq= Infinity,
    itcount = 0,
    conv    = 0,  (* number of times chisq has changed only little *)
    Ndata   = Length[data],
    Nparams = Length[paramsymbols]},
   {stddev,maxits,guess,simp,sigma,tracing} =
   {Errors,MaxIterations,InitialGuess,Simplifier,ErrorNotation,ProgressTrace}
   /.{opts}/.Options[NonlinearFit];
   If[NumberQ[stddev],
      stddev = Table[N[stddev],{Ndata}],
      If[Head[stddev]==List,
      If[Length[stddev]<Ndata, (* short stddev, pad with average *)
         stddev = Join[stddev,Table[Release[N[(Plus@@stddev)/Length[stddev]]],
                                    {Ndata-Length[stddev]} ]];
         Message[NonlinearFit::short],
         If[Length[stddev]>Ndata, (* long stddev, truncate *)
            stddev = N[Drop[stddev,Ndata-Length[stddev]]];
            Message[NonlinearFit::long]]]]];
  stddev = N[stddev];
  If[!TrueQ[Head[stddev]==List] || !EveryQ[stddev,_Real],
     Message[NonlinearFit::baddev,stddev];
     stddev = Table[Errors/.Options[NonlinearFit],{Ndata}]];
  x = (Drop[#,-1]&) /@ data; (* separate abscissas and ordinates *)
  y = Last /@ data;
  guess = N[guess];
  If[MatchQ[guess,_List] && EveryQ[guess,_Real] && Length[guess] == Nparams,
     params = guess,
     params = Table[1.,{Nparams}]; (* default guess: all unity *)
     guess == Automatic || Message[NonlinearFit::guess,guess,paramsymbols] ];
  Clear[f,Df];
  If[MatchQ[vars,_Symbol],
     f[vars1_,(Pattern[#,Blank[]]&) /@ paramsymbols] =
        (model/.vars->vars1);
     Df[vars1_,(Pattern[#,Blank[]]&) /@ paramsymbols] =
       (simp[D[model/.vars->vars1,#]]&) /@ paramsymbols,
     vars1 = (Unique[]&) /@ vars;
     f[Pattern[#,Blank[]]& /@ vars1,Pattern[#,Blank[]]& /@ paramsymbols] =
       (model/.ReplaceLists[vars,vars1]);
     Df[Pattern[#,Blank[]]& /@ vars1,Pattern[#,Blank[]]& /@ paramsymbols] =
       (simp[D[model/.ReplaceLists[vars,vars1],#]]&) /@ paramsymbols];
  While[
     itcount++ < maxits && conv < 2,
     If[tracing,
        Print[StringForm["NonlinearFit::trace params = ``, lambda = ``.", params,lambda]] ];
     {params,chisq,lambda,alpha} =
       MarquardtIteration[x,y,stddev,params,f,Df,lambda];
     If[chisq<oldChisq && (oldChisq-chisq<0.1 || oldChisq-chisq<0.001 chisq),
        conv++,
        conv=0];
     oldChisq = chisq ];
  If[conv < 2, Message[NonlinearFit::noconv]];
  covar = Inverse[alpha];
  Join[MapN[#1->#2&,paramsymbols,params],
       MapN[sigma[#1]->Sqrt[#2]&,paramsymbols,Diagonal[covar]] ]
] /;
(MatrixQ[data] && EveryQ[data,_?(EveryQ[#,_Real]&)]
   || Message[NonlinearFit::baddata,data]) &&
(EveryQ[vars,_Symbol] && Length[vars]==Max[Length[First[data]]-1,1]
   || Message[NonlinearFit::badvar,vars,Max[Length[First[data]]-1,1]]) &&
(EveryQ[paramsymbols,_Symbol] 
   || Message[NonlinearFit::badpar,paramsymbols])

NonlinearFit\
  [data_List, model_, paramsymbols_List, vars_List,opts___]:=
    nlfit[Transpose[{Range[1.,N[Length[data]]],N[data]}],model,
          paramsymbols,vars,opts]/;
  EveryQ[data,_?NumberQ]

NonlinearFit[data_List,model_,paramsymbols_List,vars_List,opts___]:=
  nlfit[N[data],model,paramsymbols,vars,opts]

NonlinearFit\
  [data_List, model_, paramsymbols_List, vars_Symbol, opts___]:=
    NonlinearFit[data,model,paramsymbols,{vars},opts]

End[]

EndPackage[]

Print[{NonlinearFit}]

Null
