Levenberg-Marquardt法あたりの実装練習

　非線形最小二乗問題を解く上で、

1次近似を使う最急降下法
2次近似を使うニュートン法
ニュートン法のヘッセ行列をヤコビ行列の積で近似するガウスニュートン法
ガウスニュートン法と最急降下法を上手く混ぜるLevenberg-Marquardt法

と整理できて、これらを実装してみる。

　サンプルの関数としてお馴染みのRosenbrock functionを使う。

　(1, 1)のところに極小値が存在する。

　これに対してPyhtonで各種最適化手法を実装した。Line Searchなどは使わず、Levenberg-Marquardt法でもλの混合具合の調整などはせず、かなりシンプルに関係を把握する目的のコードになる。

実装

import numpy as np
from copy import deepcopy
from collections import defaultdict
import matplotlib.pyplot as plt


def error1(x1: float, x2: float) -> float:
    return x1 - 1


def error2(x1: float, x2: float) -> float:
    return 10 * (x2 - x1**2)


def rosenbrock_function(x1: float, x2: float) -> float:
    return 1 / 2 * (error1(x1, x2) ** 2 + error2(x1, x2) ** 2)


def rosenbrock_gradient(x1: float, x2: float) -> np.ndarray:
    g1 = error1(x1, x2) * 1.0 + error2(x1, x2) * (-20 * x1)
    g2 = error1(x1, x2) * 0.0 + error2(x1, x2) * 10
    return np.array([g1, g2])


def hessian_matrix(x1: float, x2: float) -> np.ndarray:
    return np.array([[1 - 200 * x2 + 600 * x1**2, -200 * x1], [-200 * x1, 100]])


def jacobian_matrix(x1: float, x2: float) -> np.ndarray:
    row1 = np.array([1, 0])
    row2 = np.array([-20 * x1, 10])
    return np.vstack([row1, row2])


if __name__ == "__main__":
    np.set_printoptions(formatter={"float": "{:.3f}".format})

    initial_guess = np.array([-10.985, 50.940])

    # Newton method
    x = deepcopy(initial_guess)
    dict_newton = defaultdict(list)
    dict_newton["x1"].append(x[0])
    dict_newton["x2"].append(x[1])
    dict_newton["f(x)"].append(rosenbrock_function(x[0], x[1]))
    for i in range(100):
        g = rosenbrock_gradient(x[0], x[1])
        h = hessian_matrix(x[0], x[1])
        x -= np.linalg.inv(h) @ g
        dict_newton["x1"].append(x[0])
        dict_newton["x2"].append(x[1])
        dict_newton["f(x)"].append(rosenbrock_function(x[0], x[1]))
        if np.linalg.norm(g) < 1e-6:
            break
    plt.plot(dict_newton["x1"], dict_newton["x2"], marker="o", label="Newton")

    # Gauss-Newton method
    # print("Gauss-Newton method")
    x = deepcopy(initial_guess)
    dict_gauss_newton = defaultdict(list)
    dict_gauss_newton["x1"].append(x[0])
    dict_gauss_newton["x2"].append(x[1])
    dict_gauss_newton["f(x)"].append(rosenbrock_function(x[0], x[1]))
    for i in range(100):
        g = rosenbrock_gradient(x[0], x[1])
        j = jacobian_matrix(x[0], x[1])
        h_hat = j.T @ j
        x -= np.linalg.inv(h_hat) @ g
        dict_gauss_newton["x1"].append(x[0])
        dict_gauss_newton["x2"].append(x[1])
        dict_gauss_newton["f(x)"].append(rosenbrock_function(x[0], x[1]))
        if np.linalg.norm(g) < 1e-6:
            break
    plt.plot(dict_gauss_newton["x1"], dict_gauss_newton["x2"], marker="x", label="Gauss-Newton")

    # Gradient descent method
    x = deepcopy(initial_guess)
    dict_gradient_descent = defaultdict(list)
    dict_gradient_descent["x1"].append(x[0])
    dict_gradient_descent["x2"].append(x[1])
    dict_gradient_descent["f(x)"].append(rosenbrock_function(x[0], x[1]))
    for i in range(100):
        g = rosenbrock_gradient(x[0], x[1])
        x -= 0.000025 * g
        dict_gradient_descent["x1"].append(x[0])
        dict_gradient_descent["x2"].append(x[1])
        dict_gradient_descent["f(x)"].append(rosenbrock_function(x[0], x[1]))
        if np.linalg.norm(g) < 1e-6:
            break
    plt.plot(dict_gradient_descent["x1"], dict_gradient_descent["x2"], marker="s", label="Gradient descent")

    # Levenberg-Marquardt method
    x = deepcopy(initial_guess)
    dict_levenberg_marquardt = defaultdict(list)
    dict_levenberg_marquardt["x1"].append(x[0])
    dict_levenberg_marquardt["x2"].append(x[1])
    dict_levenberg_marquardt["f(x)"].append(rosenbrock_function(x[0], x[1]))
    for i in range(100):
        g = rosenbrock_gradient(x[0], x[1])
        j = jacobian_matrix(x[0], x[1])
        h_hat = j.T @ j
        x -= np.linalg.inv(h_hat + 0.1 * np.eye(2)) @ g
        dict_levenberg_marquardt["x1"].append(x[0])
        dict_levenberg_marquardt["x2"].append(x[1])
        dict_levenberg_marquardt["f(x)"].append(rosenbrock_function(x[0], x[1]))
        if np.linalg.norm(g) < 1e-6:
            break
    plt.plot(dict_levenberg_marquardt["x1"], dict_levenberg_marquardt["x2"], marker="^", label="Levenberg-Marquardt")

    plt.scatter(initial_guess[0], initial_guess[1], color="blue", label="Initial guess")
    plt.scatter(1, 1, color="red", label="Minimum")
    plt.legend()
    plt.savefig("optimization.png", bbox_inches="tight", pad_inches=0.05)