Graphing Scipy optimize.minimize convergence results each iteration?

Question

I would like to perform some tests on my optimisation routine using scipy.optimize.minimize, in particular graphing the convergence (or the rather objective function) each iteration, over multiple tests.

Suppose I have the following linearly constrained quadratic optimisation problem:

minimise: x_i Q_ij x_j + a|x_i|

subject to: sum(x_i) = 1

I can code this as:

def _fun(x, Q, a):
    c = np.einsum('i,ij,j->', x, Q, x)
    p = np.sum(a * np.abs(x))
    return c + p
def _constr(x):
    return np.sum(x) - 1

And I will implement the optimisation in scipy as:

x_0 = # some initial vector
x_soln = scipy.optimise.minimize(_fun, x_0, args=(Q, a), method='SLSQP', 
                                 constraints={'type': 'eq', 'fun': _constr})

I see that there is a callback argument but which only accepts a single argument of the parameter values at each iteration. How can I utilise this in a more esoteric case where I might have other arguments that need to be supplied to my callback function?

Answer 1

The way I solved this was use a generic callback cache object referenced each time from my callback function. Let's say you want to do 20 tests and plot the objective function after each iteration in the same chart. You will need an outer loop to run 20 tests, but we'll create that later.

First lets create a class that will store all iteration objective function values for us, and a couple extra bits and pieces:

class OpObj(object):
    def __init__(self, Q, a):
        self.Q, self.a = Q, a
        rv = np.random.rand()
        self.x_0 = np.array([rv, (1-rv)/2, (1-rv)/2])
        self.f = np.full(shape=(500,), fill_value=np.NaN)
        self.count = 0
    def _fun(self, x):
        return _fun(x, self.Q, self.a)

Also lets add a callback function that manipulates that class obj. Don't worry that it has more than one argument for now since we'll fix this later. Just make sure the first parameter is the solution variables.

def cb(xk, obj=None):
    obj.f[obj.count] = obj._fun(xk)
    obj.count += 1

All this does is use the object's functions and values to update itself, counting the number of iterations each time. This function will be called after each iteration.

Putting this all together all we need is two more things: 1) some matplotlib-ing to do the plot, and fixing the callback to have only one argument. We can do that with a decorator which is exactly what functools partial does. It returns a function with less arguments than the original. So the final code looks like this:

import matplotlib.pyplot as plt
import scipy.optimize as op
import numpy as np
from functools import partial

Q = np.array([[1.0, 0.75, 0.45], [0.75, 1.0, 0.60], [0.45, 0.60, 1.0]])
a = 1.0

def _fun(x, Q, a):
    c = np.einsum('i,ij,j->', x, Q, x)
    p = np.sum(a * np.abs(x))
    return c + p
def _constr(x):
    return np.sum(x) - 1

class OpObj(object):
    def __init__(self, Q, a):
        self.Q, self.a = Q, a
        rv = np.random.rand()
        self.x_0 = np.array([rv, (1-rv)/2, (1-rv)/2])
        self.f = np.full(shape=(500,), fill_value=np.NaN)
        self.count = 0
    def _fun(self, x):
        return _fun(x, self.Q, self.a)

def cb(xk, obj=None):
    obj.f[obj.count] = obj._fun(xk)
    obj.count += 1

fig, ax = plt.subplots(1,1)
x = np.linspace(1,500,500)
for test in range(20):
    op_obj = OpObj(Q, a)
    x_soln = op.minimize(_fun, op_obj.x_0, args=(Q, a), method='SLSQP',
                         constraints={'type': 'eq', 'fun': _constr},
                         callback=partial(cb, obj=op_obj))
    ax.plot(x, op_obj.f)

ax.set_ylim((1.71,1.76))
plt.show()