_LMFit
Description
_LMFit carries out the Levenberg-Marquardt iterations for NLS: starting from the initial parameter values, it adjusts the parameters and the damping until the sum of squared residuals stops decreasing, then returns the final parameters together with the information NLS needs to report convergence or failure.
Used by: NLS
Calls: _NumJac
Source Code
_LMFit = LAMBDA(
fw, // LAMBDA(p): n x 1 column of weighted fitted values
p0, // starting parameters as a column
swY, // weighted observations, sqrt(w) * Y
tol, // convergence tolerance on the relative (scaled) step
maxiter, // maximum number of iterations
LET(
k, ROWS(p0),
SS0, SUMSQ(swY - fw(p0)),
// one iteration; st is the state column described above
lmStep, LAMBDA(st,
LET(
p, TAKE(st, k),
lam, INDEX(st, k + 1),
SS, INDEX(st, k + 2),
iters, INDEX(st, k + 4),
J, _NumJac(fw, p),
IF(
NOT(ISNUMBER(SUM(J))),
VSTACK(p, lam, SS, 1, iters, 1),
LET(
r, swY - fw(p),
// J'J with its columns scaled to unit length
Dg, TOCOL(BYCOL(J ^ 2, LAMBDA(c_, SUM(c_)))),
D, IF(Dg > 0, SQRT(Dg), 1),
As, MMULT(TRANSPOSE(J), J) / MMULT(D, TRANSPOSE(D)),
gs, MMULT(TRANSPOSE(J), r) / D,
// SS after a step with damping lam_t
SSat, LAMBDA(lam_t,
SUMSQ(swY - fw(p + MMULT(MINVERSE(As + lam_t * MUNIT(k)), gs) / D))
),
// first t (0 to 9) for which damping lam*10^t reduces SS, or -1
t_ok, REDUCE(
-1,
SEQUENCE(10, 1, 0),
LAMBDA(acc, t,
IF(AND(acc >= 0), acc, IF(AND(IFERROR(SSat(lam * 10 ^ t) < SS, FALSE)), t, -1))
)
),
IF(
AND(t_ok < 0),
VSTACK(p, lam * 10 ^ 10, SS, --(lam * 10 ^ 10 >= 10 ^ 16), iters + 1, 0),
LET(
lam_t, lam * 10 ^ t_ok,
ds, MMULT(MINVERSE(As + lam_t * MUNIT(k)), gs),
delta, ds / D,
pn, p + delta,
SSn, SUMSQ(swY - fw(pn)),
conv, AND(
SQRT(SUMSQ(D * delta)) <= tol * SQRT(SUMSQ(D * pn)),
lam_t < 1
),
// Nielsen: rho = actual / predicted reduction
rho, (SS - SSn) / SUM(ds * (lam_t * ds + gs)),
VSTACK(pn, lam_t * MAX(1 / 3, 1 - (2 * rho - 1) ^ 3), SSn, --conv, iters + 1, 0)
)
)
)
)
)
),
// AND() makes each IF condition a single value: a 1x1 array
// condition (as from INDEX of a REDUCE result) makes IF evaluate
// both branches and broadcast, so a finished fit would keep iterating
REDUCE(
VSTACK(p0, 10 ^ -3, SS0, 0, 0, 0),
SEQUENCE(maxiter),
LAMBDA(st, i_, IF(AND(INDEX(st, k + 3) = 1), st, lmStep(st)))
)
)
);