Levenberg Marquardt implementation in numpy

lm_numpy.py

import sympy as sp
import numpy as np

def levenberg(sympy_func, xi, target, sympy_param, guess, weight=None, module='numpy'):
    '''
    Computes the minimum `guess` so that `sympy_func(guess) - target` is minimized
    
    Arguments:
    sympy_func  : Should be a sympy function to be minimized
    xi          : numpy array, the x-axis of `target`
    target      : numpy array, data the function should be fit to
    sympy_param : tuple of sympy symbols or string, determines the order of the parameters
    guess       : tuple of floats, initial guess for the parameters in order of sympy_param
    
    Returns:
    p           : tuple of floats, improved guess
    
    Example:
    import sympy as sp
    from sympy.abc import x, a, b, c, d
    
    f = a*x**3 + b*x**2 + c*x + d

    true_par = (2,3,-4, 3)
    func = sp.lambdify((x, a, b, c, d), f)

    xi = np.linspace(-100, 100, 1000)
    target = func(xi, *true_par)
    
    p = levenberg(f, xi, target, sympy_param=(a,b,c,d), guess=(1,1,1,1))    

    plt.figure()
    plt.plot(xi, target)
    plt.plot(xi, func(xi, *p), ls = '--')
    '''
    p = np.array(guess, like=xi)
    epsilon_1 = 1e-3
    epsilon_2 = 1e-3
    epsilon_3 = 1e-1
    epsilon_4 = 1e-1
    lam = 1e-2
    lambda_DN_fac = 9
    lambda_UP_fac = 11
    MAX_ITERATIONS = 200

    
    def chi_square(diff, weight=1):
        return diff.T @ (weight * diff[:, None])

    def lm_matx(diff, guess, weight, J):
        X2 = chi_square(diff, weight)
        JtWJ = J.T @ (J * weight)
        JtWdy = J.T @ (weight * diff[:, None])
        return X2, JtWJ, JtWdy


    if weight is not None:
        if np.var(weight) == 0: # all weights are equal
            W = weight[0] * np.ones((len(x), 1))
        else:
            W = np.abs(weight)
    else:
        W = 1
        
    DoF = len(xi) - len(p) + 1 #degrees of freedom

    p_min   = -100*abs(p)
    p_max   =  100*abs(p)
    
    #sympy_param = tuple(symb for symb in sympy_func.free_symbols if str(symb) != 'x')
    def jacobian_func(p):
        F = sp.Matrix([sympy_func])
        J = F.jacobian(sympy_param)
        lamb = sp.lambdify((x, *sympy_param), J, modules=module)
        arr = lamb(xi, *p)[0]
        return np.array([entry if np.shape(entry) == xi.shape else np.ones_like(xi)*entry for entry in arr], like=xi).T
    
    lamb = sp.lambdify(("x",) + sympy_param, sympy_func, modules=module)
    
    def func(p):
        return lamb(xi, *p)

    J = jacobian_func(p)
    diff = target - func(p)
    X2, JtWJ, JtWdy = lm_matx(diff, p, W, J)
    X2_old = X2

    iteration = 0
    while iteration < MAX_ITERATIONS:
        iteration += 1
        h = np.linalg.solve(JtWJ + lam*np.diag(np.diag(JtWJ)),JtWdy)[:,0]

        p_try = p + h
        p_try = np.minimum(np.maximum(p_min,p_try), p_max)
        fit = func(p_try)
        diff = target - fit
        X2_try = chi_square(diff, W)

        rho = (X2 - X2_try) / ( h.T @ (lam * h + JtWdy[:,0]) );

        if rho > epsilon_4:
            dX2 = X2 - X2_old
            X2_old = X2
            p_old = p
            y_old = fit
            p = p_try
            diff = target - func(p)
            J = jacobian_func(p)
            X2, JtWJ, JtWdy = lm_matx(diff, p, W, J)

            lam = max(lam/lambda_DN_fac,1e-7)

        else:
            X2 = X2_old
            diff = target - func(p)
            J = jacobian_func(p)
            X2, JtWJ, JtWdy = lm_matx(diff, p, W, J)
            lam = max(lam/lambda_UP_fac,1e7)

        # convergence conditions
        JtWdy_convergence = max(abs(JtWdy)) < epsilon_1
        parameter_convergence = np.max(np.abs(h)/(abs(p)+1e-12)) < epsilon_2
        chi2_convergence = X2/DoF < epsilon_3

        convergence = [JtWdy_convergence, parameter_convergence, chi2_convergence]

        if any(convergence)  &  (iteration > 2) :
            break
    return p