v923z / micropython-ulab

a numpy-like fast vector module for micropython, circuitpython, and their derivatives
https://micropython-ulab.readthedocs.io/en/latest
MIT License
401 stars 112 forks source link

[FEATURE REQUEST] Nelder-mead multiparametric optimisation #560

Open furbrain opened 1 year ago

furbrain commented 1 year ago

Describe the solution you'd like I'd like to be able to do a multiparametric optimisation, ideally using the nelder-mead algorithm,as in scipy https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html.

I have rewritten the python code currently used in scipy, to be compatible with ulab (below) but there doesn't seem to be a way to include python code with ulab

Additional context This is some calibration code for a magnetometer. This is tweaking the final results by adding a radial basis function to the output of the magnetometer output, which can account for non-linearity in the sensor. I have this working on my laptop, but I'd like to be able to use it on an embedded device

"""
This is an implementation of the nelder-mead optimisation algorithm.
"""

try:
    import numpy as np
    import warnings
except ImportError:
    from ulab import numpy as np

def minimize_neldermead(func, x0, args=(), maxiter=None, disp=False,
                        xatol=1e-4, fatol=1e-4,):
    """
    Minimization of scalar function of one or more variables using the
    Nelder-Mead algorithm.
    Options
    -------
    disp : bool
        Set to True to print convergence messages.
    maxiter: int
        Maximum allowed number of iterations.
        Will default to ``N*200``, where ``N`` is the number of
        variables.
    xatol : float, optional
        Absolute error in xopt between iterations that is acceptable for
        convergence.
    fatol : number, optional
        Absolute error in func(xopt) between iterations that is acceptable for
        convergence.
    References
    ----------
    .. [1] Gao, F. and Han, L.
       Implementing the Nelder-Mead simplex algorithm with adaptive
       parameters. 2012. Computational Optimization and Applications.
       51:1, pp. 259-277
    """

    rho = 1
    chi = 2
    psi = 0.5
    sigma = 0.5

    nonzdelt = 0.05
    zdelt = 0.00025

    N = len(x0)

    # create the initial simplex
    sim = np.empty((N + 1, N), dtype=x0.dtype)
    sim[0] = x0
    for k in range(N):
        y = x0[:]
        if y[k] != 0:
            y[k] = (1 + nonzdelt) * y[k]
        else:
            y[k] = zdelt
        sim[k + 1] = y

    # If neither are set, then set both to default
    if maxiter is None :
        maxiter = N * 200

    one2np1 = list(range(1, N + 1))
    fsim = np.full((N + 1,), np.inf)

    for k in range(N + 1):
        fsim[k] = func(sim[k])
    iterations = 1
    while iterations < maxiter:
        order = np.argsort(fsim,axis=0)
        best = order[0]
        worst = order[-1]
        second_worst = order[-2]
        sim_best = sim[best]
        fsim_best = fsim[best]
        sim_worst = sim[worst]
        fsim_worst = fsim[worst]
        if ((np.max(abs(sim - sim_best))) <= xatol and
                np.max(abs(fsim - fsim_best)) <= fatol):
            break

        # calculate centroid, by calculating sum of all vertices, minus the worst
        xbar = (np.sum(sim,axis=0) - sim_worst) / N
        xr = (1 + rho) * xbar - rho * sim_worst
        fxr = func(xr)
        doshrink = 0

        if fxr < fsim[0]:
            xe = (1 + rho * chi) * xbar - rho * chi * sim_worst
            fxe = func(xe)

            if fxe < fxr:
                sim[worst] = xe
                fsim[worst] = fxe
            else:
                sim[worst] = xr
                fsim[worst] = fxr
        else:  # fsim[0] <= fxr
            if fxr < fsim[second_worst]:
                sim[worst] = xr
                fsim[worst] = fxr
            else:  # fxr >= fsim[-2]
                # Perform contraction
                if fxr < fsim_worst:
                    xc = (1 + psi * rho) * xbar - psi * rho * sim_worst
                    fxc = func(xc)

                    if fxc <= fxr:
                        sim[worst] = xc
                        fsim[worst] = fxc
                    else:
                        doshrink = 1
                else:
                    # Perform an inside contraction
                    xcc = (1 - psi) * xbar + psi * sim_worst
                    fxcc = func(xcc)

                    if fxcc < fsim_worst:
                        sim[worst] = xcc
                        fsim[worst] = fxcc
                    else:
                        doshrink = 1

                if doshrink:
                    for j in one2np1:
                        sim[j] = sim_best + sigma * (sim[j] - sim_best)
                        fsim[j] = func(sim[j])
        iterations += 1

    x = sim[0]
    fval = np.min(fsim)

    if iterations >= maxiter:
        msg = "Max Iterations"
        if disp:
            warnings.warn(msg, RuntimeWarning, 3)
    else:
        msg = "Success"
        if disp:
            print(msg)
            print("         Current function value: %f" % fval)
            print("         Iterations: %d" % iterations)

    return {"status": msg,
            "x": x,
            "simplex": sim,
            "fsim": fsim,
            "iterations": iterations}

class CheckOptimize:
    """ Base test case for a simple constrained entropy maximization problem
    (the machine translation example of Berger et al in
    Computational Linguistics, vol 22, num 1, pp 39--72, 1996.)
    """

    def setup_method(self):
        self.F = np.array([[1, 1, 1],
                           [1, 1, 0],
                           [1, 0, 1],
                           [1, 0, 0],
                           [1, 0, 0]])
        self.K = np.array([1., 0.3, 0.5])
        self.startparams = np.zeros(3)
        self.solution = np.array([0., -0.524869316, 0.487525860])
        self.maxiter = 1000
        self.funccalls = 0
        self.gradcalls = 0
        self.trace = []

    def func(self, x):
        self.funccalls += 1
        if self.funccalls > 6000:
            raise RuntimeError("too many iterations in optimization routine")
        log_pdot = np.dot(self.F, x)
        logZ = np.log(sum(np.exp(log_pdot)))
        f = logZ - np.dot(self.K, x)
        self.trace.append(x[:])
        return f

chk = CheckOptimize()
chk.setup_method()
res = minimize_neldermead(chk.func, np.array((1.0, 1.0, 1.0)), disp=True)
print(res)
print(chk.funccalls)
v923z commented 1 year ago

You can actually add python stuff as demonstrated in https://github.com/v923z/micropython-ulab/tree/master/snippets. I think this would be a brilliant addition, if you care to create a PR.

It might make sense to re-implement at least part of this in C. A lot is happening here, and I have the feeling that this is probably expensive.

furbrain commented 1 year ago

I can certainly create a PR, but the stuff in snippets doesn't seem to be being built into ulab firmware - am I missing something?

v923z commented 1 year ago

No, you are not missing anything. The idea behind the snippets was that people could easily add extra functions without having to recompile the firmware. But the functions are still added to the name space.

Having said that, I think that this function should be implemented in C. The nested loops are expensive in python. I actually entertained the idea of adding the simplex method, but there seemed to be no demand for that, so I dropped it.

By the way, I wonder, whether

sim = np.empty((N + 1, N), dtype=x0.dtype)

is correct: what happens, if you start out with an integer type?

furbrain commented 1 year ago

I don't really speak Micropython flavoured C, so I think I might struggle to translate this function. I'd be happy to tidy it up and add it as a snippet PR however - would you like me to do so?

I've just copied this from the scipy implementation, and then converted to use the ulab functions, so I don't really know what would happen with an integer type.

v923z commented 1 year ago

I don't really speak Micropython flavoured C, so I think I might struggle to translate this function. I'd be happy to tidy it up and add it as a snippet PR however - would you like me to do so?

Yes, I think it would be great. It would also ease a bit on the time pressure. I'm willing to do this properly in C, but if there was a not-so-performant python implementation in the meantime, that would be very useful.

I've just copied this from the scipy implementation, and then converted to use the ulab functions, so I don't really know what would happen with an integer type.

I haven't tried it, either, but in the contraction step you are bound to leave the set of integers, I think.