How to write objective functions

optimagic is very flexible when it comes to the objective function and its derivatives. In this how-to guide we start with simple examples, that would also work with scipy.optimize before we show advanced options and their advantages.

The simplest case

In the simplest case, fun maps a numpy array into a scalar objective value. The name of first argument of fun is arbitrary.

import numpy as np
import optimagic as om


def sphere(x):
    return x @ x


res = om.minimize(
    fun=sphere,
    params=np.arange(3),
    algorithm="scipy_lbfgsb",
)
res.params.round(6)

array([ 0., -0.,  0.])

More flexible params

In all but the most simple problems, a flat numpy array is not ideal to keep track of all the different parameters one wants to optimize over. Therefore, optimagic accepts objective functions that work with other parameter formats. Below we show a simple example. More examples can be found here.

def dict_fun(x):
    return x["a"] ** 2 + x["b"] ** 4


res = om.minimize(
    fun=dict_fun,
    params={"a": 1, "b": 2},
    algorithm="scipy_lbfgsb",
)

res.params

{'a': 2.2389496983474372e-07, 'b': 0.005724249755832909}

The important thing is that the params provided to minimize need to have the format that is expected by the objective function.

Functions with additional arguments

In many applications, the objective function takes more than params as argument. This can be achieved via fun_kwargs. Take the following simplified example:

def shifted_sphere(x, offset):
    return (x - offset) @ (x - offset)


res = om.minimize(
    fun=shifted_sphere,
    params=np.arange(3),
    algorithm="scipy_lbfgsb",
    fun_kwargs={"offset": np.ones(3)},
)
res.params

array([0.99999999, 0.99999999, 0.99999999])

fun_kwargs is a dictionary with keyword arguments for fun. There is no constraint on the number or names of those arguments.

Least-Squares problems

Many estimation problems have a least-squares structure. If so, specialized optimizers that exploit this structure can be much faster than standard optimizers. The sphere function from above is the simplest possible least-squarse problem you could imagine: the least-squares residuals are just the params.

To use least-squares optimizers in optimagic, you need to mark your function with a decorator and return the least-squares residuals instead of the aggregated function value.

@om.mark.least_squares
def ls_sphere(params):
    return params


res = om.minimize(
    fun=ls_sphere,
    params=np.arange(3),
    algorithm="pounders",
)
res.params.round(5)

array([ 0.,  0., -0.])

Any least-squares optimization problem is also a standard optimization problem. You can therefore optimize least-squares functions with scalar optimizers as well:

res = om.minimize(
    fun=ls_sphere,
    params=np.arange(3),
    algorithm="scipy_lbfgsb",
)
res.params.round(5)

array([ 0., -0.,  0.])

Returning additional information

You can return additional information such as intermediate results, debugging information, etc. in your objective function. This information will be stored in a database if you use logging.

To do so, you need to return a FunctionValue object.

def sphere_with_info(x):
    return om.FunctionValue(value=x @ x, info={"avg": x.mean()})


res = om.minimize(
    fun=sphere_with_info,
    params=np.arange(3),
    algorithm="scipy_lbfgsb",
)

res.params.round(6)

array([ 0., -0.,  0.])

The info can be an arbitrary dictionary. In the oversimplified example we returned the mean of the parameters, which could have been recovered from the params history that is collected anyways but in real applications this feature can be helpful.