Nelder-Mead Method

Copyright (C) 2022 François Chollet, derived from this code
Copyright (C) 2022 Andreas Kloeckner

Licensed under the GNU General Public License Version 2

Resources:

• Nelder-Mead method (Wikipedia)

import numpy as np
import matplotlib.pyplot as plt
import copy

This is a challenge problem for optimization algorithms known as Rosenbrock's banana function.

def f(X):
    x = X[0]
    y = X[1]
    val = 100.0 * (y - x**2)**2 + (1.0 - x)**2
    return val

fig = plt.figure(figsize=(6, 6))

xmesh, ymesh = np.mgrid[-2:2:50j,-2:2:50j]
fmesh = f(np.array([xmesh, ymesh]))

plt.axis("equal")
plt.contour(xmesh, ymesh, fmesh, 150)

<matplotlib.contour.QuadContourSet at 0x7f702f0b8b80>

Set the starting guess:

x_start = np.array([2, 2./5])

Set some parameters:

step = 0.1
alpha = 1.
gamma = 2.
rho = -0.5
sigma = 0.5

Some initialization:

dim = len(x_start)
prev_best = f(x_start)
no_improv = 0
res = [[x_start, prev_best]]

for i in range(dim):
    x = copy.copy(x_start)
    x[i] = x[i] + step
    score = f(x)
    res.append([x, score])

And the iteration:

def iteration(res):
    res = res[:]
    
    # order
    res.sort(key=lambda x: x[1])
    best = res[0][1]

    # centroid
    x0 = [0.] * dim
    for tup in res[:-1]:
        for i, c in enumerate(tup[0]):
            x0[i] += c / (len(res)-1)

    # reflection
    xr = x0 + alpha*(x0 - res[-1][0])
    rscore = f(xr)
    if res[0][1] <= rscore < res[-2][1]:
        del res[-1]
        res.append([xr, rscore])
        return res

    # expansion
    if rscore < res[0][1]:
        xe = x0 + gamma*(x0 - res[-1][0])
        escore = f(xe)
        if escore < rscore:
            del res[-1]
            res.append([xe, escore])
            return res
        else:
            del res[-1]
            res.append([xr, rscore])
            return res

    # contraction
    xc = x0 + rho*(x0 - res[-1][0])
    cscore = f(xc)
    if cscore < res[-1][1]:
        del res[-1]
        res.append([xc, cscore])
        return res

    # reduction
    x1 = res[0][0]
    nres = []
    for tup in res:
        redx = x1 + sigma*(tup[0] - x1)
        score = f(redx)
        nres.append([redx, score])
    return nres

res = iteration(res)

fig = plt.figure(figsize=(9, 9))

xmesh, ymesh = np.mgrid[-2:2:50j,-2:2:50j]
fmesh = f(np.array([xmesh, ymesh]))

plt.axis("equal")
plt.contour(xmesh, ymesh, fmesh, 150)

pts = np.array([pt for pt, val in res]).T
plt.plot(pts[0], pts[1], "ko-")

[<matplotlib.lines.Line2D at 0x7f702e80e2c0>]