Cleaned up Levenberg-Marquardt a little.

1 files changed, 42 insertions, 42 deletions

diff --git a/schofield.wl b/schofield.wl
index 1c750a9..eb0525c 100644
--- a/schofield.wl
+++ b/schofield.wl
@@ -187,7 +187,7 @@ eq[n_, g_][m_, p_, q_] := Flatten[Join[{ruleθ0[g], ruleAH[g], g'[0] θc - 1}, e
 
 eqAround[n_, g_][m_, p_, q_] := Flatten[Join[{ruleθ0[g], ruleAH[g]}, eqLow[n, g][#] & /@ Range[0, m],eqMid[RF[n], g][#] & /@ Range[0, p], eqHigh[n, g] /@ Range[2, q, 2]]] //. rules[g]
 
-formResiduals[data_, functions_, δ_:0] := If[Head[#1]===Around,
+formResiduals[data_, functions_, δ_:10^(-15)] := If[Head[#1]===Around,
   (#1["Value"] - #2) / Max[#1["Uncertainty"], δ],
   (#1 - #2) / δ] & @@@ Thread[{data, functions}]
 
@@ -208,54 +208,54 @@ newSol[eqs_, oldSol_, newVars_, δ_:0, γ_:0, opts___] := FindRoot[
   opts
 ]
 
-levenburgMarquardt[r_, \[Beta]0_, \[Lambda]0_ : 1, \[Nu]_ : 2] := 
+levenbergMarquardtStep[M_, λ_, g_] := LinearSolve[M + λ DiagonalMatrix[Diagonal[M]], g]
+
+levenburgMarquardt[r_, β0_, λ0_ : 1, ν_ : 2] := 
  Module[
-    {
-     n, \[Beta], new\[Beta], rc, J, I, oldJ, newJ, oldr, newr, M, 
-   g, \[Delta], newcost, \[Lambda], oldcost, temp
+  {
+    n = Length[β0],
+    β = β0,
+    λ = λ0,
+    newβ = β0,
+    x, rf, Jf, oldJ, oldr, newr, M, g, δ, oldC, newC
    },
-  n = Length[\[Beta]0]; \[Lambda] = \[Lambda]0; \[Beta] = \[Beta]0;
   PrintTemporary["Compiling the Jacobian..."];
-  rc = Compile[{{x, _Real, 1}}, 
-      Evaluate[
-       r /. Thread[Rule[First /@ \[Beta], Part[x, #] & /@ Range[n]]]]];
-  J = Compile[{{x, _Real, 1}}, 
-      Evaluate[
-       D[r, {First /@ \[Beta]}] /. 
-      Thread[Rule[First /@ \[Beta], Part[x, #] & /@ Range[n]]]]];
+  rf = Compile[{{x, _Real, 1}}, 
+    Evaluate[r /. Thread[Rule[First /@ β, Part[x, #] & /@ Range[n]]]]
+  ];
+  Jf = Compile[{{x, _Real, 1}}, 
+    Evaluate[D[r, {First /@ β}] /. Thread[Rule[First /@ β, Part[x, #] & /@ Range[n]]]]
+  ];
   PrintTemporary["Beginning the algorithm."];
-  oldJ = J[\[Beta][[All, 2]]];
-  oldr = rc[\[Beta][[All, 2]]];
-  oldcost = Re[Total[oldr^2]];
-  new\[Beta] = \[Beta];
+  oldr = rf[β[[All, 2]]];
+  oldJ = Jf[β[[All, 2]]];
+  oldC = Re[Total[oldr^2]];
   g = Re[Transpose[oldJ] . oldr];
   M = Re[Transpose[oldJ] . oldJ];
-  \[Delta] = LinearSolve[M + \[Lambda] DiagonalMatrix[Diagonal[M]], g];
-  PrintTemporary[Dynamic[oldcost]]
-  While[Norm[\[Delta]] > 10^-15,
-   \[Delta] = 
-    LinearSolve[M + \[Lambda]/\[Nu] DiagonalMatrix[Diagonal[M]], g];
-   new\[Beta][[All, 2]] -= \[Delta];
-   newr = rc[new\[Beta][[All, 2]]];
-   newcost = Re[Total[newr^2]];
-   While[newcost > oldcost,
-      \[Delta] = 
-     LinearSolve[M + \[Lambda] DiagonalMatrix[Diagonal[M]], g];
-    new\[Beta] = \[Beta];
-    new\[Beta][[All, 2]] -= \[Delta];
-    newr = rc[new\[Beta][[All, 2]]];
-    newcost = Re[Total[newr^2]];
-    \[Lambda] *= \[Nu];
+  δ = levenbergMarquardtStep[M, λ / ν, g];
+  PrintTemporary[Dynamic[oldC]]
+  While[Norm[δ] > 10^-15,
+    newβ[[All, 2]] -= δ;
+    newr = rf[newβ[[All, 2]]];
+    newC = Re[Total[newr^2]];
+    While[newC > oldC,
+      δ = levenbergMarquardtStep[M, λ, g];
+      newβ = β;
+      newβ[[All, 2]] -= δ;
+      newr = rf[newβ[[All, 2]]];
+      newC = Re[Total[newr^2]];
+      λ *= ν;
     ];
-   \[Lambda] /= \[Nu];
-   oldcost = newcost;
-   oldr = newr;
-   \[Beta] = new\[Beta];
-   oldJ = J[\[Beta][[All, 2]]];
-   g = Re[Transpose[oldJ] . oldr];
-   M = Re[Transpose[oldJ] . oldJ];
-   ];
-  {newcost, \[Beta]}
+    λ /= ν;
+    oldC = newC;
+    oldr = newr;
+    β = newβ;
+    oldJ = Jf[β[[All, 2]]];
+    g = Re[Transpose[oldJ] . oldr];
+    M = Re[Transpose[oldJ] . oldJ];
+    δ = levenbergMarquardtStep[M, λ / ν, g];
+  ];
+  {newC, β}
  ]
 
 EndPackage[]