scipy.optimize.minimize, computing the constraints and their jacobian in one go

Question

There is quite a bit of shared computation between the constraints I have on a minimization problem and their Jacobian, to the point where I get the Jacobian almost for free. Is there any way to share the computation?

Answer 1

Since constraints and jacobians are probably not always evaluated together, you can only expect a smaller improvement. But if you can put the common calculations into a separate function/method, you can cache its return values so that you don't need to recompute them later:

import scipy.optimize as opt
from functools import lru_cache

# only for the dummy example:
import scipy as sp
from time import sleep

def cost(x):
    '''dummy cost function to minimize on [1,11]'''
    return sp.log(x)

@lru_cache(maxsize=None) # careful with this choice
def heavy_stuff(x):
    '''idle computation representing common work in constraint and jacobian'''
    sleep(0.1)
    return 0

def constraint(x):
    '''constraint for [1,11], with simulated idle work'''
    # the following only works for 1d arrays, needs more work for nd
    throwaway = heavy_stuff(tuple(x))  
    return 5 - abs(6 - x)  # non-negative between 1 and 11

def jacobian(x):
    '''return the jacobiam with the same simulated idle work'''
    throwaway = heavy_stuff(tuple(x))
    return 1/x

x0 = 11
tol = 0
opts = {'disp': True}
cons = {'type': 'ineq', 'fun': constraint}
kwargs = {'method':'SLSQP', 'constraints': cons,
          'jac': jacobian, 'tol': tol, 'options': opts}
res = opt.minimize(cost,x0,**kwargs)
print(heavy_stuff.cache_info())

The dummy example above attempts to minimize log(x) over the interval [1,11]. Instead of constant bounds I defined a constraint that gives us the interval, so that I can show what I mean regarding your question.

Both constraint and jacobian do the same work, this is what you want to spare in case multiple evaluations happen with the same argument. You have to put all these common calculations into a common function (here named heavy_stuff), and work with the return values in both constraint and jacobian.

My point is that you should use functools.lru_cache to memoize the heavy stuff. By setting an appropriate cache size, multiple evaluations of heavy_stuff with the same x will give you the previously computed return value at once, without having to redo the calculation.

If my suspicion is correct, a maxsize=1 would probably suffice inside the lru_cache decorator. Setting maxsize=None (no upper limit) poses the danger of losing too much memory for no good reason. You should experiment and see if multiple memoized values are necessary, or whether a few or just one is enough.

Note, however, that lru_cache uses a dict to look up previously computed results, in which keys are the input parameters of the decorated function. This means that the input arguments have to be hashable, which practically means that they should be immutable. Numpy arrays are very similar to lists, and they are similarly not hashable. This is why we call heavy_stuff with tuple(x): the 1d array input is converted to a tuple. If x is a multidimensional array, then every level of nesting has to be converted to a tuple. What's worse, heavy_stuff almost certainly has to convert the tuples back to numpy ndarrays in order to do the heavy lifting. However, if computing the jacobian/constraint is really this CPU-intensive, then you're probably still better off overall.

If you want to assess the performance of your cache, you should take a closer look at heavy_stuff.cache_info() printed at the end. It will tell you how often cached values were used, and how many times new values had to be calculated.