Parameter Calibration

from abc import abstractmethod
from collections.abc import Callable
from typing import Self

import numpy as np
from matplotlib import pyplot as plt
from scipy import spatial
from scipy.integrate import solve_ivp
from scipy.optimize import minimize
from scipy.stats import qmc

class LithiumClusterConfig:
    def __init__(self) -> None:
        """Class for configuring a LithiumClusterModel
        :return:
        """
        self.kr: float = 0.0
        self.kf: float = 0.0
        self.dr: float = 0.0
        self.dm: float = 0.0
        self.lf: float = 0.0
        self.pc: float = 0.0
        self.tstop: float = 0.0
        self.f0 = 0.0
        self.m0 = 0.0
        self.r0 = 0.0

    def from_array(self, params: np.ndarray) -> None:
        """Initialises the configuration with a specific parameter and input selection
        :params: list or array containing the parameter/input values
        :return:
        """
        if np.any(params) < 0:
            print("values smaller than zero detected..truncating")
            params = np.array([max(p, 0) for p in params])
        self.kr: float = params[0]
        self.kf: float = params[1]
        self.dr: float = params[2]
        self.dm: float = params[3]
        self.lf: float = params[4]
        self.pc: float = params[5]
        self.tstop: float = params[6]
        self.f0: float = params[7]
        self.m0: float = params[8]
        self.r0: float = params[9]

    def to_array(self) -> np.ndarray:
        """Converts the configuration object to an array
        :return:
        """
        return np.array(
            [
                self.kr,
                self.kf,
                self.dr,
                self.dm,
                self.lf,
                self.pc,
                self.tstop,
                self.f0,
                self.m0,
                self.r0,
            ]
        )

    def copy(self) -> Self:
        """Deep copy of the configuration object
        :return: copy
        """
        s = LithiumClusterConfig()
        s.from_array(self.to_array())
        return s

    def __str__(self) -> str:
        """String representation
        :return: a summary of the model parameters
        """
        return ",".join(["{:6s}".format(f"{x:.02f}") for x in self.to_array()])


class LithiumClusterModel(LithiumClusterConfig):
    def __init__(self, config: LithiumClusterConfig) -> None:
        """Initialises the model with a specific parameter selection
        :param config: instance of the configuration object
        :return:
        """
        super().__init__()
        self.from_array(config.to_array())

    def p(self, t) -> float:
        """Bombardment function
        :param t: current time
        :return: current bombardment impact
        """
        return self.pc if t <= self.tstop else 0.0

    def rhs(self, t: float, x: np.ndarray) -> np.ndarray:
        """Right hand side of the ODE model
        :param t: t
        :param x: x
        :return: dx/dt
        """
        p = self.p(t)
        f = x[0]
        m = x[1]
        r = x[2]
        df = (
            p
            - self.lf * f
            - 2 * self.kf * f**2
            - self.kr * f * m
            + 2 * self.dm * m
            + self.dr * r
        )
        dm = self.kf * f**2 - self.kr * f * m - self.dm * m + self.dr * r
        dr = self.kr * f * m - self.dr * r
        return np.array([df, dm, dr])

    def run(self, tend: float) -> tuple[np.ndarray, np.ndarray]:
        """Executes the model over a certain timespan
        :param tend: simulation time-span
        :return: array with time values, 2d array with F-,M-,R- center concentration values
        """
        y0 = np.array([self.f0, self.m0, self.r0])
        sol = solve_ivp(self.rhs, y0=y0, t_span=[0, tend], method="BDF")
        return sol.t, sol.y

# simple evaluatuon with default parameters
default_params = LithiumClusterConfig()
default_params.kr = 1.0
default_params.kf = 0.1
default_params.lf = 1000.0
default_params.dr = 0.1
default_params.dm = 1.0
default_params.pc = 10000.0
default_params.tstop = 10.0
default_params.f0 = 9.975
default_params.m0 = 1.674
default_params.r0 = 84.99
model = LithiumClusterModel(default_params)
T, X = model.run(20)
plt.figure(figsize=(8, 4))
plt.plot(T, X[0, :], label="f(t)")
plt.plot(T, X[1, :], label="m(t)")
plt.plot(T, X[2, :], label="r(t)")
plt.legend()
plt.show()

# parameter variation for lf
plt.figure(figsize=(8, 4))
for lf in [10 ** ((5 + i) / 3) for i in range(1, 8)]:
    params = default_params.copy()
    params.lf = lf
    model = LithiumClusterModel(params)
    T, X = model.run(10)
    plt.plot(T, X[0, :], label=f"f(t;{lf:.0f})")
plt.gca().set_yscale("log")
plt.gca().set_xscale("log")
plt.legend()
plt.show()

n_data = 200
np.random.seed(12345)  # fix the seed to be redproducible
model = LithiumClusterModel(default_params)
t_end = 20.0
t_sol, y_sol = model.run(t_end)  # execute the model with the default parameters
# create a random time-base with n_data points
t_ref = t_end * np.random.random(n_data)
t_ref.sort()
# interpolate the sim-result and add noise
y_ref = np.zeros([3, n_data])
y_ref[0, :] = np.interp(t_ref, t_sol, y_sol[0, :]) + 0.5 * np.random.randn(n_data)
y_ref[1, :] = np.interp(t_ref, t_sol, y_sol[1, :]) + 0.5 * np.random.randn(n_data)
y_ref[2, :] = np.interp(t_ref, t_sol, y_sol[2, :]) + 2 * np.random.randn(n_data)
# compare sim-result with synthetic data
plt.figure(figsize=(12, 5))
for i, lbl in enumerate(["f(t)", "m(t)", "r(t)"]):
    plt.subplot(1, 3, i + 1)
    plt.plot(t_sol, y_sol[i, :], "k", label="sim_result")
    plt.plot(t_ref, y_ref[i, :], "ob", label="synthetic data")
    plt.legend()
    plt.title(lbl)
plt.show()

def error_lithium_cluster(t_sim: np.ndarray, y_sim: np.ndarray) -> float:
    """Computes the error between the simulation results and the reference data
    :param t_sim: vector with time-values for the simulation result
    :param y_sim: matrix with f,m,r curves for the simulation result
    :return: error between simulation and reference data
    """
    t_cmp = np.arange(0, 20, 0.2)
    err = [0, 0, 0]
    for i in range(3):
        y_ref_cmp = np.interp(t_cmp, t_ref, y_ref[i, :])
        y_ref_sim = np.interp(t_cmp, t_sim, y_sim[i, :])
        err[i] = ((y_ref_cmp - y_ref_sim) * (y_ref_cmp - y_ref_sim)).sum() / 100.0
    return sum(err)


model = LithiumClusterModel(
    default_params
)  # execute the model with the default parameters
t_end = 20.0
t_sol, y_sol = model.run(t_end)
print(f"Reference error: {error_lithium_cluster(t_sol, y_sol)}")

full_parameter_space = np.array(
    [
        [0.2, 3.0],
        [0.05, 0.2],
        [0.01, 0.2],
        [0.3, 3.0],
        [100.0, 10000.0],
        [2000.0, 20000.0],
        [5.0, 20.0],
        [0.0, 10.0],
        [0.0, 100.0],
        [0.0, 20.0],
    ]
)  # all boundaries of all paramaters


def extract_free_parameters(
    free_parameter_indices: np.array, full_parameter_vector: np.array
) -> np.array:
    """Method to extract the subvector of free parameters from a full parameter vector
    :param free_parameter_indices: indices to free parameters
    :param full_parameter_vector: full parameter vector
    :return: subvector containing the free parameter values
    """
    return full_parameter_vector[free_parameter_indices]


def extract_free_parameter_space(free_parameter_indices: np.array):
    """Returns the parameter space to the free parameters
    :param free_parameter_indices: indices to free parameters
    :return: space for the free parameters
    """
    return full_parameter_space[free_parameter_indices, :]


def embed_free_parameters(
    free_parameter_indices: np.array, free_parameter_vector: np.array
) -> np.array:
    """Embeds the vector of free parameters into a vector of full parameters assuming that all other parameter values are known from the default parametrisation
    :param free_parameter_indices: indices to free parameters
    :param free_parameter_vector: vector with free parameter values
    :return: full vector with embedded values
    """
    arr = default_params.to_array()
    arr[free_parameter_indices] = free_parameter_vector
    return arr

class Optimizer:
    def __init__(
        self, function: Callable[[np.ndarray], float] | None, name: str
    ) -> None:
        """Base class for optimization of methods
        :param function: function to minimize. Can be None to set it later on.
        :param name: name of the method
        """
        self.name = name
        if function == None:
            self.function_base = lambda x: 0
        else:
            self.function_base = function
        self.count = 0
        self.function_base = function
        self.p_sol = None
        self.err_sol = None

    def set_function(self, function: Callable[[np.ndarray], float]) -> None:
        """Can set the function also a-posteriori
        :param function: objective function to be minimized
        """
        self.function_base = function
        self.reset()

    def set_ground_truth(self, p_sol: np.ndarray) -> None:
        """Defines the reference solution to evaluate the performance of the method
        :param p_sol: reference parameter set
        """
        self.p_sol = p_sol
        self.err_sol = self.function_base(self.p_sol)

    def reset(self) -> None:
        """Resets the internal counter which tracks function evaluations"""
        self.count = 0

    def function(self, p: np.ndarray) -> float:
        """Applies the funtion to a given parameter vector and increments the counter
        :param p: parameter vector
        :return: value
        """
        self.count += 1
        return self.function_base(p)

    @abstractmethod
    def run(self, max_iter: int, p0: np.ndarray = None, color="r") -> np.ndarray:
        """Abstract run method.
        Performs some iteration until (at most) max_iter and uses, optionally, a starting value p0
        :param max_iter: maximum iterations
        :param p0: starting value
        :param color: optional color specification for plotting
        :return: optimized parameter vector
        """
        return np.array([])

    def stop(self) -> bool:
        """Override this method to specify a stopping criterium. Called after every update-step
        (Currently unused)
        :return: true, if the method should stop at this point
        """
        return False


class IterativeOptimizer(Optimizer):
    def __init__(
        self, function: Callable[[np.ndarray], float] | None, name: str
    ) -> None:
        """Wrapper class for every optimized which iteratively improves one parameter set at a time.
        :param function: function to minimize. Can be None to speficy later
        :param name: name of the method
        """
        super().__init__(function, name)

    def reset(self) -> None:
        """Resets all internal variables related to covergence tracking and output visualization"""
        super().reset()
        self.trajectory = list()
        self.fps = list()
        self.step = 0

    @abstractmethod
    def update_step(self, p: np.ndarray, fp: float) -> tuple[np.ndarray, float | None]:
        """Abstract method to performs an update step.
        :param p: current status of the iteration
        :param fp: function value at p
        :return: updated state and function value at this point. For the latter, return None, if it needs to be computed anew.
        """
        return np.array([]), None

    def print_msg(self) -> None:
        """Prints a message to analyze the convergence status at the current step
        :return:
        """
        if (
            len(self.fps) == 1 or self.fps[-1] != self.fps[-2]
        ):  # only print a message if we have some progress to show
            print(f"step {self.step}: {self.trajectory[-1]} -> {self.fps[-1]}")

    def print_final_msg(self) -> None:
        """Prints a summary message after the iteration
        :return:
        """
        if self.p_sol is not None:
            diff = self.trajectory[-1] - self.p_sol
            reldiff = diff / self.p_sol
            accx = np.linalg.norm(reldiff)
        else:
            accx = "undef."
        print(
            f"Summary: steps {self.step}, func-evals: {self.count}, accuracy x: {accx}, residual: {self.function_base(self.trajectory[-1])}"
        )

    def run(self, max_iter: int, p0: np.ndarray = None, color="r") -> np.ndarray:
        """Runs the iteration until (at most) max_iter (depending on the implementation of "stop").
        :param max_iter: maximum iterations
        :param p0: initial guess (must be given in this implementation!)
        :param color: color for plotting
        :return: optimized parameter set
        """
        assert p0 is not None, ValueError(
            "Iterative optimizer must have a p0 value which is not None"
        )
        self.reset()
        p = p0
        fp = self.function(p)
        self.trajectory.append(p)
        self.fps.append(fp)
        # print state for the initial condition
        self.print_msg()
        # start iteration
        for i in range(max_iter):
            self.step = i + 1
            p, fp = self.update_step(p, fp)
            self.trajectory.append(p)
            if fp == None:
                fp = self.function(p)
            self.fps.append(fp)
            self.print_msg()
            if self.stop():
                break
        # plot the trajectory
        # label for plot
        lbl = f"{self.name} (in {self.step} steps ({self.count} F-evals) from {self.fps[0]:.02f} to {self.fps[-1]:.02f})"
        n = len(self.trajectory)
        for i, (x0, x1) in enumerate(zip(self.trajectory[:-1], self.trajectory[1:])):
            plt.plot(
                [x0[0], x1[0]],
                [x0[1], x1[1]],
                color=color,
                alpha=(i + 1 + n / 5) / (n - 1 + n / 5),
                label=lbl if i == n - 2 else None,
            )
        self.print_final_msg()
        return self.trajectory[-1]


class NewtonOptimizer(IterativeOptimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float] | None,
        h: float = 0.01,
        gamma: float = 0.5,
    ) -> None:
        """Implements the Newton method
        :param function: function to minimize. Can be None to specify later
        :param h: stepsize for estimation of the gradient
        :param gamma: stepsize for damped Newton step
        """
        super().__init__(function, name="newton")
        self.h = h
        self.gamma = gamma

    def update_step(self, p: np.ndarray, fp: float) -> tuple[np.ndarray, float | None]:
        """Performs a Newton step
        :param p: current status of the iteration
        :param fp: function value at p
        :return: updated state and corresponding function value (will be always None in this method)
        """
        dim = len(p)
        # compute gradient and hessian
        id = np.identity(dim)
        grad = np.zeros(dim)
        hess = np.zeros([dim, dim])
        fpi = list()
        # evaluate the error function at selected points in proximity
        for i in range(dim):
            fpi.append(self.function(p + self.h * id[:, i]))
        for i in range(dim):
            grad[i] = (fpi[i] - fp) / self.h  # compute the gradient
            for j in range(dim):
                fpij = self.function(p + self.h * id[:, i] + self.h * id[:, j])
                hess[i, j] = (
                    fpij - fpi[i] - fpi[j] + fp
                ) / self.h**2  # compute the hessian
        # use linear-solve instead of explicitly computing the inverse
        diff = np.linalg.solve(hess, self.gamma * grad)
        # update
        return p - diff, None


class GradientDescentOptimizer(IterativeOptimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float] | None,
        h: float = 0.01,
        alpha: float = 0.1,
    ) -> None:
        """Implements the Gradient Descent method
        :param function: function to minimize. Can be None to specify later
        :param h: stepsize for estimation of the gradient
        :param alpha: stepsize for the descent step
        """
        super().__init__(function, name="gradient-descent")
        self.h = h
        self.alpha = alpha

    def update_step(self, p: np.ndarray, fp: float) -> tuple[np.ndarray, float | None]:
        """Performs a Gradient Descent step
        :param p: current status of the iteration
        :param fp: function value at p
        :return: updated state and corresponding function value (will be always None in this method)
        """
        dim = len(p)
        id = np.identity(dim)
        grad = np.zeros(dim)
        # compute the gradient
        for i in range(dim):
            grad[i] = (self.function(p + self.h * id[:, i]) - fp) / self.h
        # update
        return p - self.alpha * grad, None


def compute_parameter_space_heatmap(
    function: Callable[[np.ndarray], float],
    x_min: float,
    x_max: float,
    x_steps: int,
    y_min: float,
    y_max: float,
    y_steps: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Method to compute a 2D heatmap (x and y) of (a part of) the (first two dimensons of the) parameterspace.
    I.e. it computes f(x,y) for a grid of x and y values and returns it as a matrix
    :param function: function to evaluate
    :param x_min: minimum x-value for the grid
    :param x_max: maximum x-value for the grid
    :param x_steps: number of x-steps. I.e the x-vector will be x_steps elements long
    :param y_min: minimum y-value for the grid
    :param y_max: maximum y-value for the grid
    :param y_steps: number of y-steps. I.e the y-vector will be x_steps elements long
    :return: tuple with three elements: (1) vector of x-values, (2) vector of y-values, (3) matrix of function values, whereas the rows correspond to the y-values and the columns to the x-values
    """
    dx = (x_max - x_min) / x_steps
    dy = (y_max - y_min) / y_steps
    xs = np.arange(x_min, x_max + dx, dx)
    ys = np.arange(y_min, y_max + dy, dy)
    matrix = np.zeros([len(ys), len(xs)])
    for i, x in enumerate(xs):
        for j, y in enumerate(ys):
            matrix[j, i] = function(np.array([x, y]))
    return xs, ys, matrix


def plot_parameter_space_heatmap(xs: np.ndarray, ys: np.ndarray, zs: np.ndarray):
    """Supplements the method compute_parameter_space_heatmap plotting the heatmap and contour lines.
    I.e. the output of the method can be used directly as input to this one.
    The reason why these two methods are separated is, because a computed heatmap can be left in the workspace this way and must not be recomputed all the time whenever it should be plotted.
    :param xs: vector of x-values
    :param ys: vector of y-values
    :param zs: matrix of function values, whereas the rows correspond to the y-values and the columns to the x-values
    :return:
    """
    zs2 = (
        zs**0.25
    )  # take the 4-th root to better visualise differences for small values
    dx = xs[1] - xs[0]
    dy = ys[1] - ys[0]
    plt.imshow(
        zs2,
        extent=(xs[0] - dx / 2, xs[-1] + dx / 2, ys[0] - dy / 2, ys[-1] + dy / 2),
        cmap="viridis",
        alpha=0.3,
        origin="lower",
    )  # setting extent and origin is crucial for putting everything at the right place
    X, Y = np.meshgrid(xs, ys)  # meshgrid required for contour plot
    plt.contour(X, Y, zs2, levels=30, linewidths=0.5, colors=[0.5, 0.5, 0.5])


function = lambda x: (
    (x[0] + 2 * x[1] - 7) ** 2 + (2 * x[0] + x[1] - 5) ** 2
)  # booth function
p_sol = np.array([1, 3])  # actual solution
p0 = np.array([0.5, 1.5])  # starting value for the iteration

# compute a heatmap around the solution to create a nice visualisation
plt.figure(figsize=[8, 8])
plot_parameter_space_heatmap(
    *compute_parameter_space_heatmap(
        function, p_sol[0] - 2, p_sol[0] + 2, 100, p_sol[1] - 2, p_sol[1] + 2, 100
    )
)

# apply the Newton method
opt = NewtonOptimizer(function, h=0.01, gamma=0.5)
print(f"##### {opt.name.upper()} ####")
opt.set_ground_truth(p_sol)
opt.run(4, p0=p0, color="r")

# apply the Gradient descent
opt = GradientDescentOptimizer(function, h=0.01, alpha=0.08)
print(f"##### {opt.name.upper()} ####")
opt.set_ground_truth(p_sol)
opt.run(10, p0=p0, color="b")

plt.legend()
plt.show()

def visualise_lithium_simulation_fit(
    free_parameters: np.ndarray, free_parameter_indices: np.ndarray
) -> None:
    """Visualizes the esults of the Lithium-Cluster model with the fitted and the reference parameters.
    :param free_parameters: fitted parameters
    :param free_parameter_indices: indices of the free parameters
    """
    params = default_params.copy()
    model = LithiumClusterModel(params)
    T, X = model.run(20)  # run the model with the reference parameters

    p2arr = embed_free_parameters(free_parameter_indices, free_parameters)
    params2 = LithiumClusterConfig()
    params2.from_array(p2arr)
    model2 = LithiumClusterModel(params2)
    T2, X2 = model2.run(20)  # run the model with the fitted parameters

    plt.plot(T, X[0, :], color="r", linestyle="dashed", label="f(t) (ref)")
    plt.plot(T2, X2[0, :], color="r", label="f(t) (fit)")
    plt.plot(T, X[1, :], color="b", linestyle="dashed", label="m(t) (ref)")
    plt.plot(T2, X2[1, :], color="b", label="m(t) (fit)")
    plt.plot(T, X[2, :], color="g", linestyle="dashed", label="r(t) (ref)")
    plt.plot(T2, X2[2, :], color="g", label="r(t) (fit)")


lithium_cluster_heatmap_2d = None


def run_lithium_cluster_test(
    optimizer: Optimizer | list[Optimizer],
    free_parameter_indices: np.ndarray,
    max_iters: int | list[int],
):
    global lithium_cluster_heatmap_2d

    def function(p: np.ndarray):  # wrapped lithium cluster model
        params = default_params.to_array()
        params[free_parameter_indices] = p
        config = LithiumClusterConfig()
        config.from_array(params)
        mdl = LithiumClusterModel(config)
        t_sim, x_sim = mdl.run(20)
        return error_lithium_cluster(t_sim, x_sim)

    if isinstance(optimizer, Optimizer):
        optimizer = [optimizer]
    if isinstance(max_iters, int):
        max_iters = [max_iters for opt in optimizer]

    p_sol = extract_free_parameters(
        free_parameter_indices, default_params.to_array()
    )  # (usually unknown) target solution
    free_parameter_space = extract_free_parameter_space(free_parameter_indices)

    p0 = 0.5 * np.sum(
        free_parameter_space, 1
    )  # pick starting value for iteration right in the middle of the parameter space

    plt.figure(figsize=(10, 5))
    if (
        len(free_parameter_indices) == 2
        and free_parameter_indices[0] == 0
        and free_parameter_indices[1] == 1
    ):
        ## if we solve a 2d problem, plot a heatmap
        if lithium_cluster_heatmap_2d is None:
            print("compute parameter space heatmap (might take a while)")
            lithium_cluster_heatmap_2d = compute_parameter_space_heatmap(
                function, 0.5, 2.0, 25, 0.05, 0.15, 25
            )
        plot_parameter_space_heatmap(*lithium_cluster_heatmap_2d)
    plt.scatter(
        [p_sol[0]], [p_sol[1]], color="k", marker="x", label="reference minimum"
    )  # plot the actual global optimum
    colors = ["r", "b", "g", "c", "m", "y"]
    sols = list()
    for i, (opt, m_it) in enumerate(zip(optimizer, max_iters)):
        opt.set_function(function)
        opt.set_ground_truth(p_sol)
        print(
            f"##### Start fitting of {len(free_parameter_indices)} dimensional problem with {opt.name.upper()} ####"
        )
        sol = opt.run(
            m_it, p0=p0, color=colors[i % 6]
        )  # we should not test more than 6 at once anyway ...
        sols.append(sol)
    plt.legend()
    plt.axis("auto")
    plt.title(f"fit of {len(free_parameter_indices)} dimensional problem")
    plt.figure(figsize=(10, 5))
    for i, (sol, opt) in enumerate(zip(sols, optimizer)):
        plt.subplot(1, len(sols), i + 1)
        visualise_lithium_simulation_fit(sol, free_parameter_indices)
        plt.legend()
        plt.title(f"fitted with {opt.name}")


newton_opt = NewtonOptimizer(None, h=0.001, gamma=0.5)
gradient_opt = GradientDescentOptimizer(None, h=0.001, alpha=0.0000002)
run_lithium_cluster_test([newton_opt, gradient_opt], np.array([0, 1]), [15, 40])
plt.show()

class SciPyOptimizer(IterativeOptimizer):
    def __init__(
        self, function: Callable[[np.ndarray], float] | None, method_name: str
    ) -> None:
        """Implements the Gradient Descent method
        :param function: function to minimize. Can be None to specify later
        :param h: stepsize for estimation of the gradient
        :param alpha: stepsize for the descent step
        """
        super().__init__(function, name="scipy-" + method_name)
        self.method_name = method_name

    def run(self, max_iter: int, p0: np.ndarray = None, color="r") -> np.ndarray:
        """Runs the scipy optimizer. Max-iter will be ignored.
        :param max_iter: maximum iterations (will be ignored)
        :param p0: initial guess (must be given in this implementation!)
        :param color: color for plotting
        :return: optimized parameter set
        """
        assert p0 is not None, ValueError(
            "Iterative optimizer must have a p0 value which is not None"
        )
        self.reset()
        p = p0.copy()
        fp = self.function_base(p)  # dont count this function evaluation
        self.trajectory.append(p)
        self.fps.append(fp)
        self.print_msg()
        popt = minimize(self.function, x0=p, method=self.method_name)
        p = popt.x
        fp = self.function_base(p)  # dont count this function evaluation
        self.step = popt.nit
        self.trajectory.append(p)
        self.fps.append(fp)
        self.print_msg()
        lbl = f"{self.name} (in {self.step} steps ({self.count} F-evals) from {self.fps[0]:.02f} to {self.fps[-1]:.02f})"
        for i, (x0, x1) in enumerate(zip(self.trajectory[:-1], self.trajectory[1:])):
            plt.plot([x0[0], x1[0]], [x0[1], x1[1]], color=color, alpha=1.0, label=lbl)
        self.print_final_msg()
        return self.trajectory[-1]


bfgs_opt = SciPyOptimizer(None, "BFGS")
cg_opt = SciPyOptimizer(None, "CG")
run_lithium_cluster_test([bfgs_opt, cg_opt], np.array([0, 1]), [0, 0])
plt.show()

newton_opt = NewtonOptimizer(None, h=0.001, gamma=0.5)
gradient_opt = GradientDescentOptimizer(None, h=0.001, alpha=0.0000003)
run_lithium_cluster_test(
    [newton_opt, gradient_opt], np.array([0, 1, 2, 3, 6]), [15, 40]
)
plt.show()

class HillClimbingOptimizer(IterativeOptimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float] | None,
        parameter_space: np.ndarray,
        alpha=0.1,
    ) -> None:
        """Optimizer using a stochastic Hill Climbing algorithm
        :param function: function to minimize. Can be None to specify later
        :param parameter_space: boundaries of the parameter space
        :param alpha: "stepwidth", i.e. relative standarddeviation of the normal distribution.
        """
        super().__init__(function, name="hill-climber")
        self.parameter_space_lb = parameter_space[:, 0]
        self.parameter_space_up = parameter_space[:, 1]
        self.parameter_diff = self.parameter_space_up - self.parameter_space_lb
        self.alpha = alpha

    def map_to_parameterspace(self, p: np.ndarray) -> np.ndarray:
        """Maps a parameter vector (back) onto the parameterspace
        :param p: parameter vector
        :return: parameter vector mapped onto the parameterspace
        """
        return np.minimum(
            self.parameter_space_up, np.maximum(self.parameter_space_lb, p)
        )

    def mutate(self, p: np.ndarray) -> np.ndarray:
        """Performs a stochastic mutation of a parameter vector
        :param p: parameter vector
        :return: new mutated parameter vector
        """
        dim = len(p)
        X = np.random.randn(dim)
        alpha = self.alpha * 0.99**self.step
        return p + X * self.parameter_diff * alpha

    def update_step(self, p: np.ndarray, fp: float) -> tuple[np.ndarray, float | None]:
        """Performs a stocahstic hill climber step.
        Randomly mutates the parameter vector, maps it back onto the parameter space and accepts the mutation if it produces a lower error as the original.
        :param p: current status of the iteration
        :param fp: function value at p
        :return: updated state
        """
        p2 = self.mutate(p)
        p2 = self.map_to_parameterspace(p2)
        fp2 = self.function(p2)
        if fp2 < fp:
            return p2, fp2
        return p, fp


np.random.seed(12346)  # fix random seed (has quites om influence, try to vary it!)
free_parameter_indices = np.array([0, 1])
hc_opt = HillClimbingOptimizer(
    None,
    alpha=0.1,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
)
run_lithium_cluster_test(hc_opt, free_parameter_indices, 100)
free_parameter_indices = np.array([0, 1, 2, 3, 6])
hc_opt = HillClimbingOptimizer(
    None,
    alpha=0.1,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
)
run_lithium_cluster_test(hc_opt, free_parameter_indices, 200)
plt.show()

class SimulatedAnnealingOptimizer(HillClimbingOptimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float] | None,
        parameter_space: np.ndarray,
        alpha: float = 0.1,
        t0: float = 10.0,
        qt: float = 0.95,
    ) -> None:
        """Optimizer using a Simulated Annealing
        :param function: function to minimize. Can be None to specify later
        :param parameter_space: boundaries of the parameter space
        :param alpha: "stepwidth" for mutation, i.e. relative standarddeviation of the normal distribution.
        :param t0: initial temperature
        :param qt: decline of temperature per step
        """
        super().__init__(function, parameter_space, alpha)
        self.name = "simulated-annealing"
        self.t0 = t0
        self.qt = qt

    def update_step(self, p: np.ndarray, fp: float) -> tuple[np.ndarray, float | None]:
        """:param p: current status of the iteration
        :param fp: function value at p
        :return: updated state
        """
        p2 = self.mutate(p)
        p2 = self.map_to_parameterspace(p2)
        fp2 = self.function(p2)
        if fp2 < fp:
            return p2, fp2
        t = self.t0 * self.qt**self.step
        if np.random.random() < np.exp(-(fp2 - fp) / t):
            return p2, fp2
        return p, fp


np.random.seed(12346)
free_parameter_indices = np.array([0, 1])
hc_opt = SimulatedAnnealingOptimizer(
    None,
    alpha=0.1,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
    t0=10,
    qt=0.95,
)
run_lithium_cluster_test(hc_opt, free_parameter_indices, 100)
free_parameter_indices = np.array([0, 1, 2, 3, 6])
hc_opt = SimulatedAnnealingOptimizer(
    None,
    alpha=0.1,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
    t0=10,
    qt=0.95,
)
run_lithium_cluster_test(hc_opt, free_parameter_indices, 200)
plt.show()

class PopulationBasedOptimizer(Optimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float] | None,
        name: str,
        parameter_space: np.ndarray,
    ) -> None:
        """Base class for a population-based optimization metaheuristic.
        :param function: function to minimize, can be None and speficied later
        :param name: name of the method
        :param parameter_space: for all population-based methods we consider parameter bounds
        """
        super().__init__(function, name)
        self.segments: dict[int, list[tuple[np.ndarray, np.ndarray]]] = (
            dict()
        )  # for plotting
        self.population: list[np.ndarray] = list()
        self.fvals: list[float] = list()
        self.min_fvals: list[float] = list()
        self.step = 0
        self.parameter_space = parameter_space
        self.parameter_space_lb = parameter_space[:, 0]
        self.parameter_space_up = parameter_space[:, 1]
        self.parameter_diff = self.parameter_space_up - self.parameter_space_lb

    def map_to_parameterspace(self, p: np.ndarray) -> np.ndarray:
        """Default method to map a parameter vector back into the parameter space
        :param p: parameter vector draft
        :return: parameter vector mapped into bounds
        """
        return np.minimum(
            self.parameter_space_up, np.maximum(self.parameter_space_lb, p)
        )

    def reset(self) -> None:
        """Resets the counter
        :return:
        """
        super().reset()
        self.segments = dict()
        self.population = list()
        self.fvals = list()
        self.min_fvals = list()
        self.step = 0

    @abstractmethod
    def create_population(self) -> tuple[list[np.ndarray], list[float | None]]:
        """Abstract method to initialize the population
        :return: list of parameter vectors and corresponding list of function values (can be None)
        """
        return list(), list()

    @abstractmethod
    def update_population(
        self, pop: list[np.ndarray], fpop: list[float]
    ) -> tuple[list[np.ndarray], list[float | None]]:
        """Abstract method to performs an update step for the population
        :param pop: current population, i.e. list of parameter vectors
        :param fpop: corresponding list of function values
        :return: updated list of parameter vectors and corresponding list of function values (can be None)
        """
        return list(), list()

    def print_msg(self) -> None:
        """Prints a message to analyze the convergence status at the current step
        :return:
        """
        mn = self.min_fvals[-1]
        i = self.fvals.index(mn)
        print(
            f"step {self.step}: best individual is {self.population[i]} with min error {mn}, max error {max(self.fvals)}"
        )

    def print_final_msg(self) -> None:
        """Prints a summary message after the iteration
        :return:
        """
        mn = self.min_fvals[-1]
        i = self.fvals.index(mn)
        opt = self.population[i]
        if self.p_sol is not None:
            diff = opt - self.p_sol
            reldiff = diff / self.p_sol
            accx = np.linalg.norm(reldiff)
        else:
            accx = "undef."
        print(
            f"Summary: steps {self.step}, func-evals: {self.count}, accuracy x: {accx}, residual: {self.function_base(opt)}"
        )

    def compute_missing_fvals(
        self, pop: list[np.ndarray], fpop: list[float | None]
    ) -> list[float]:
        """Computes all not yet computed function values to the current population
        :param pop: current population, i.e. list of parameter vectors
        :param fpop: corresponding list of function values with potential None entires
        :return: complete list of function values
        """
        fun = lambda x: x[1] if x[1] is not None else self.function(x[0])
        fpop2 = list(map(fun, zip(pop, fpop)))  # this could be parallelised nicely!
        return fpop2

    def run(self, max_iter: int, p0: np.ndarray = None, color="r") -> np.ndarray:
        """Runs the iteration until (at most) max_iter (depending on the implementation of "stop").
        :param max_iter: maximum iterations
        :param p0: initial guess (will be ignored here)
        :param color: color for plotting
        :return: optimized parameter set
        """
        self.reset()
        # initialise population
        self.population, errs = self.create_population()
        self.fvals = self.compute_missing_fvals(self.population, errs)
        self.min_fvals = [min(self.fvals)]

        # print initial conditions
        self.print_msg()
        # start iteration
        for i in range(max_iter):
            self.step = i + 1
            self.segments[self.step] = list()
            self.population, fvs = self.update_population(self.population, self.fvals)
            self.fvals = self.compute_missing_fvals(self.population, fvs)
            self.min_fvals.append(min(self.fvals))
            self.print_msg()
            if self.stop():
                break
        mn = min(self.fvals)
        imn = self.fvals.index(mn)
        pmn = self.population[imn]

        lbl = f"{self.name} (in {i + 1} steps ({self.count} F-evals) from {self.min_fvals[0]:.02f} to {self.min_fvals[-1]:.02f})"
        n = max(self.segments.keys())
        for i, segs in self.segments.items():
            xs = list()
            ys = list()
            for s1, s2 in segs:
                xs.extend([s1[0], s2[0], None])
                ys.extend([s1[1], s2[1], None])
            plt.plot(
                xs,
                ys,
                color=color,
                alpha=i / n,
                linewidth=i / n,
                label=lbl if i == n else None,
            )
        self.print_final_msg()
        return pmn


class GeneticAlgorithmOptimizer(PopulationBasedOptimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float] | None,
        parameter_space: np.ndarray,
        pop_size: int = 50,
        crossover_size: int = 10,
        mutate_size: int = 10,
        pm: float = 0.5,
        alpha: float = 0.01,
    ) -> None:
        """Optimizer using a simple genetic algorithm
        :param function: callable function to minimize. Can be None and specified later
        :param parameter_space: boudaries of the parameter space
        :param pop_size: size of the population
        :param crossover_size: number of individuals to crossover
        :param mutate_size: number of individuals to mutate
        :param pm: probaility to mutate a parameter vector entry when the individual is selected for mutation
        :param alpha: "stepsize" for mutation, i.e. standarddeviation of the normal distribution relative to the size of the parameter space
        """
        super().__init__(function, "genetic_algorithm", parameter_space)
        self.pop_size = pop_size
        self.mutate_size = mutate_size
        self.crossover_size = crossover_size
        self.select_size = self.mutate_size + self.crossover_size
        self.alpha = alpha
        self.pm = pm

    def create_population(self) -> tuple[list[np.ndarray], list[float | None]]:
        """Creates an initial population using the n-dim Sobol sequence
        :return: initial population and list of corresponding function values (always None)
        """
        dim = len(self.parameter_space_lb)
        sampler = qmc.Sobol(d=dim, scramble=False)
        sample = sampler.random_base2(int(np.log2(self.pop_size) + 1))
        pop = [
            self.parameter_space_lb + sample[i, :] * self.parameter_diff
            for i in range(self.pop_size)
        ]
        return pop, [None for p in pop]

    def crossover(self, p1: np.ndarray, p2: np.ndarray):
        """Performs a crossover between p1 and p2
        :param p1: parent individual 1
        :param p2: parent individual 2
        :return: child individual
        """
        p_new = p1.copy()
        X = np.random.random(len(p1)) < 0.5
        p_new[X] = p2[X]
        self.segments[self.step].append((p1, p_new))
        self.segments[self.step].append((p2, p_new))
        return p_new

    def mutate(self, p: np.ndarray):
        """Mutates a certain individual and returns a mutation (in bounds)
        :param p: base individual
        :return: mutated individual
        """
        dim = len(p)
        p_new = p.copy()
        for i in range(dim):
            if np.random.random() < self.pm:
                X = np.random.randn()
                alpha = self.alpha  # *0.99**self.step
                p_new[i] += X * self.parameter_diff[i] * alpha
        p_new = self.map_to_parameterspace(p_new)
        self.segments[self.step].append((p, p_new))
        return p_new

    def update_population(
        self, pop: list[np.ndarray], fpop: list[float]
    ) -> tuple[list[np.ndarray], list[float | None]]:
        """Performs a full update step for the population (selection, crossover, mutation)
        :param pop: current population
        :param fpop: corresponding function values
        :return: updated population and correspondng function values (None for the new crossovers and mutants)
        """
        pop_new = list()
        fpop_new = list()
        # join the two and sort
        pop_and_f = [(p, fp) for p, fp in zip(pop, fpop)]
        pop_and_f.sort(key=lambda x: x[1])
        # perform selection process
        pop_and_f = pop_and_f[: -self.select_size]
        pop_new.extend([x[0] for x in pop_and_f])
        fpop_new.extend([x[1] for x in pop_and_f])
        # perform crossover process
        index_list = list(range(len(pop_and_f)))
        for it in range(self.crossover_size):
            i1, i2 = np.random.choice(
                index_list, 2, replace=False
            )  # draw two different indices
            pop_new.append(self.crossover(pop_and_f[i1][0], pop_and_f[i2][0]))
            fpop_new.append(None)
        # perform mutation process
        for it in range(self.mutate_size):
            i = np.random.choice(index_list)  # draw a random index
            pop_new.append(self.mutate(pop_and_f[i][0]))
            fpop_new.append(None)
        return pop_new, fpop_new


np.random.seed(12346)
free_parameter_indices = np.array([0, 1])
ga_opt = GeneticAlgorithmOptimizer(
    None,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
    pop_size=30,
    mutate_size=5,
    crossover_size=5,
    alpha=0.03,
)
run_lithium_cluster_test(ga_opt, free_parameter_indices, 20)
free_parameter_indices = np.array([0, 1, 2, 3, 6])
ga_opt = GeneticAlgorithmOptimizer(
    None,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
    pop_size=50,
    mutate_size=10,
    crossover_size=7,
    alpha=0.03,
)
run_lithium_cluster_test(ga_opt, free_parameter_indices, 30)
plt.show()

class ParticleSwarmOptimizer(PopulationBasedOptimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float] | None,
        parameter_space: np.ndarray,
        pop_size: int = 50,
        omega: float = 0.1,
        phi: float = 0.1,
        psi: float = 0.1,
    ) -> None:
        """Optimizer using Particle Swarm Optimization
        :param function: callable function to minimize. Can be None and specified later
        :param parameter_space: boudaries of the parameter space
        :param pop_size: size of the population
        :param omega: inertia parameter for the current velocity
        :param phi: attraction to the personal optimum
        :param psi: attraction to the swarm optimum
        """
        super().__init__(function, "particle_swarm_optmisation", parameter_space)
        self.pop_size = pop_size
        self.omega = omega
        self.phi = phi
        self.psi = psi
        self.velocities: list[np.ndarray] = None

    def create_population(self) -> tuple[list[np.ndarray], list[float | None]]:
        """Creates the initial population using the Sobol sequence.
        Initializes the velocities with zero-vectors
        :return: created population and corresponding function values (all None)
        """
        dim = len(self.parameter_space_lb)
        sampler = qmc.Sobol(d=dim, scramble=False)
        sample = sampler.random_base2(int(np.log2(self.pop_size) + 1))
        pop = [
            self.parameter_space_lb + sample[i, :] * self.parameter_diff
            for i in range(self.pop_size)
        ]
        self.velocities = [np.zeros(dim) for k in range(self.pop_size)]
        self.personal_opt = [(None, None) for k in range(self.pop_size)]
        self.swarm_opt = (None, None)
        return pop, [None for p in pop]

    def update_population(
        self, pop: list[np.ndarray], fpop: list[float]
    ) -> tuple[list[np.ndarray], list[float | None]]:
        """Performs a full update step for the population.
        I.e. updates the velocities and adds them to the population.
        :param pop: current population
        :param fpop: corresponding function values
        :return: updated population and correspondng function values (None for all in this algorithm)
        """
        for k, (p, fp) in enumerate(zip(pop, fpop)):
            if self.personal_opt[k][0] is None or fp < self.personal_opt[k][1]:
                self.personal_opt[k] = (p, fp)
                if self.swarm_opt[0] is None or fp < self.swarm_opt[1]:
                    self.swarm_opt = (p, fp)
        # update velocity and positions
        pop_new = list()
        for k, p in enumerate(pop):
            x1 = self.omega * self.velocities[k]
            x2 = self.phi * (self.personal_opt[k][0] - p)
            x3 = self.psi * (self.swarm_opt[0] - p)
            self.velocities[k] = x1 + np.random.random() * x2 + np.random.random() * x3
            p_new = self.map_to_parameterspace(p + self.velocities[k])
            self.segments[self.step].append((p, p_new))
            pop_new.append(p_new)
        return pop_new, [None for x in pop_new]


np.random.seed(12346)
free_parameter_indices = np.array([0, 1])
pso_opt = ParticleSwarmOptimizer(
    None,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
    pop_size=20,
    omega=0.7,
    phi=0.4,
    psi=0.4,
)
run_lithium_cluster_test(pso_opt, free_parameter_indices, 15)
free_parameter_indices = np.array([0, 1, 2, 3, 6])
pso_opt = ParticleSwarmOptimizer(
    None,
    parameter_space=extract_free_parameter_space(free_parameter_indices),
    pop_size=30,
    omega=0.7,
    phi=0.5,
    psi=0.5,
)
run_lithium_cluster_test(pso_opt, free_parameter_indices, 30)
plt.show()

class DonwhillSimplexOptimizer(PopulationBasedOptimizer):
    def __init__(
        self,
        function: Callable[[np.ndarray], float],
        parameter_space: np.ndarray,
        alpha=1.0,
        beta=0.5,
        gamma=2.0,
        rho=-0.5,
        sigma=0.5,
    ) -> None:
        """Implementation of the Downhill Simplex Method by Nelder and Mead
        :param function: function to minimize, can be None to be set later on
        :param parameter_space: parameter boundaries
        :param alpha: reflected point parameter
        :param beta: inwards contracted point parameter
        :param gamma: expanded point parameter
        :param rho: outwards contracted point parameter
        :param sigma: shrinking parameter
        """
        super().__init__(function, "downhill_simplex", parameter_space)
        self.pop_size = len(self.parameter_space_lb) + 1
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.rho = rho
        self.sigma = sigma

    def add_simplex_segments(self, pop: list[np.ndarray]) -> None:
        """Helper method to plot the simplices.
        The probleme here is that in n-D with n>2 the projecion of a simpex into the 2-D plane is not a triangle. So we use the convex hull.
        :param pop: current simplex vetices
        :return:
        """
        points = np.array([x[:2] for x in pop])
        hull = spatial.ConvexHull(points)
        for simplex in hull.simplices:
            self.segments[self.step].append(
                (points[simplex[0], :], points[simplex[1], :])
            )

    def create_population(self) -> tuple[list[np.ndarray], list[float | None]]:
        """Creates the initial simplex which spans nicely over the parameter space.
        We use N+1 of the 2^N vertices of the 2^N corners of the N-D hyper-cube [lb+0.3(ub-lb),lb+0.7*(ub-lb)] whereas lb and ub refer to the boundaries of the parameter-space.
        :return: initial vertices and corresponding function values (will all be None here)
        """
        pop = list()
        low = self.parameter_space_lb + 0.3 * (
            self.parameter_space_up - self.parameter_space_lb
        )
        upp = self.parameter_space_lb + 0.7 * (
            self.parameter_space_up - self.parameter_space_lb
        )
        for i in range(self.pop_size - 1):
            p = low.copy()
            p[i] = upp[i]
            pop.append(p)
        pop.append(low.copy())
        return pop, [None for p in pop]

    def shrink(self, pop_sorted: list[np.ndarray]) -> list[np.ndarray]:
        """Shrinks the simplex towards its first vertex.
        I.e. the resulting list will have the same first entry as the input, but all others are shifted towards it.
        :param pop_sorted: sorted list of vertices
        :return: list of vertices shrunk to the first vertex
        """
        pop_new = [pop_sorted[0]]
        for p in pop_sorted[1:]:
            pop_new.append(p + self.sigma * (pop_sorted[0] - p))
        return pop_new

    def update_population(
        self, pop: list[np.ndarray], fpop: list[float]
    ) -> tuple[list[np.ndarray], list[float | None]]:
        """Updates the vertices of the simplex according to the rules by Nelder and Mead
        :param pop: current vertices
        :param fpop: corresponding function values
        :return: updated vertices and corresponding function values (depending on the update, either one or pop_size-1 of those are None)
        """
        # join the two and sort
        pop_and_f = [(p, fp) for p, fp in zip(pop, fpop)]
        pop_and_f.sort(key=lambda x: x[1])
        # split them up again
        pop_new = [x[0] for x in pop_and_f]
        fpop_new = [x[1] for x in pop_and_f]
        # compute candidates and start with if-else stuff
        p0 = 1 / (self.pop_size - 1) * sum(pop_new[:-1])
        pr = p0 + self.alpha * (pop_new[-1] - p0)  # reflected point
        fr = self.function(pr)
        if fr >= fpop_new[-1]:  # we made it worse..
            pci = p0 + self.beta * (pop_new[-1] - p0)  # inwards contracted point
            fci = self.function(pci)
            if fci > fpop_new[-1]:
                pop_new = self.shrink(pop_new)
                fpop_new = [fpop_new[0], *[None for x in range(self.pop_size - 1)]]
            else:
                pop_new[-1] = pci
                fpop_new[-1] = fci
        elif fr >= fpop_new[-2]:  # we made it a little bit better..
            pco = p0 + self.rho * (pop_new[-1] - p0)
            fco = self.function(pco)  # outwards contracted point
            if fco > fr:
                pop_new = self.shrink(pop_new)
                fpop_new = [fpop_new[0], *[None for x in range(self.pop_size - 1)]]
            else:
                pop_new[-1] = pco
                fpop_new[-1] = fco
        elif fr >= fpop_new[0]:  # we made it a lot better..
            pop_new[-1] = pr
            fpop_new[-1] = fr
        else:  # there might be even more potential...
            pe = p0 + self.gamma * (pop_new[-1] - p0)  # expanded point
            fe = self.function(pe)
            if fe > fr:
                pop_new[-1] = pr
                fpop_new[-1] = fr
            else:
                pop_new[-1] = pe
                fpop_new[-1] = fe
        self.add_simplex_segments(pop_new)
        return pop_new, fpop_new


np.random.seed(12346)
free_parameter_indices = np.array([0, 1])
simplex_opt = DonwhillSimplexOptimizer(
    None, parameter_space=extract_free_parameter_space(free_parameter_indices)
)
run_lithium_cluster_test(simplex_opt, free_parameter_indices, 15)
free_parameter_indices = np.array([0, 1, 2, 3, 6])
simplex_opt = DonwhillSimplexOptimizer(
    None, parameter_space=extract_free_parameter_space(free_parameter_indices)
)
run_lithium_cluster_test(simplex_opt, free_parameter_indices, 30)
plt.show()

bfgs_opt = SciPyOptimizer(None, "Nelder-Mead")
cg_opt = SciPyOptimizer(None, "Powell")
run_lithium_cluster_test([bfgs_opt, cg_opt], np.array([0, 1]), [0, 0])
run_lithium_cluster_test([bfgs_opt, cg_opt], np.array([0, 1, 2, 3, 6]), [0, 0])
plt.show()