Skip to content

_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)))
        )
    )
);