Write new objective functions#

One of the goals of the new inversion framework is to be more modular and allow users to write their own objective functions and easily plug them into the inversion framework. To do so, we need to create our own class for our new objective function and make it a child of base.Objective. Then, we’ll need to write implementations for its __call__, gradient, and hessian methods with the same signatures as the ones in base.Objective, and its n_params property. Optionally, we can also add an __init__ method to our new class to be able to pass the required arguments for it to function.

For example, let’s consider the following objective function that simply evaluates the norm of the model values:

\[\phi_\text{m}(\mathbf{m}) = \sum\limits_{i=1}^M |m_i|^2 = \lVert \mathbf{m} \rVert^2 = \mathbf{m}^\text{T} \mathbf{m}\]

Its gradient will be given by:

\[\nabla \phi_\text{m}(\mathbf{m}) = 2 \mathbf{m}\]

And its Hessian it’s the identity matrix:

\[\nabla^2 \phi_\text{m}(\mathbf{m}) = 2 \mathbb{I}\]

We can implement this objective function as follows:

import numpy as np

import inversion_ideas as ii


class ModelNorm(ii.base.Objective):
    """
    Simple model norm regularization.

    Parameters
    ----------
    n_params : int
        Number of elements expected in the ``model`` array.
    """

    def __init__(self, n_params):
        self._n_params = n_params
        self._name = "norm"

    @property
    def n_params(self):
        return self._n_params

    def __call__(self, model):
        return model.T @ model

    def gradient(self, model):
        return 2 * model

    def hessian(self, model):
        return 2 * np.eye(self.n_params)

We can instantiate this function and use its methods:

n_params = 5
model_norm = ModelNorm(n_params)
print(model_norm)
φnorm(m)
model = np.arange(n_params, dtype=np.float64)
model_norm(model)
np.float64(30.0)
model_norm.gradient(model)
array([0., 2., 4., 6., 8.])
model_norm.hessian(model)
array([[2., 0., 0., 0., 0.],
       [0., 2., 0., 0., 0.],
       [0., 0., 2., 0., 0.],
       [0., 0., 0., 2., 0.],
       [0., 0., 0., 0., 2.]])

We can also choose to change its default name through its _name attribute:

class ModelNorm(ii.base.Objective):
    """
    Simple model norm regularization.

    Parameters
    ----------
    n_params : int
        Number of elements expected in the ``model`` array.
    """

    def __init__(self, n_params):
        self._n_params = n_params
        self._name = "norm"

    @property
    def n_params(self):
        return self._n_params

    def __call__(self, model):
        return model.T @ model

    def gradient(self, model):
        return 2 * model

    def hessian(self, model):
        return 2 * np.eye(self.n_params)
model_norm = ModelNorm(n_params)
print(model_norm)
φnorm(m)

Optimizing the Hessian matrix#

In the last implementation of ModelNorm we choose to return a dense matrix in the hessian method:

model_norm.hessian(model)
array([[2., 0., 0., 0., 0.],
       [0., 2., 0., 0., 0.],
       [0., 0., 2., 0., 0.],
       [0., 0., 0., 2., 0.],
       [0., 0., 0., 0., 2.]])

If n_params is too large, then this matrix can grow too big to fit in memory. Since it’s mostly filled with zeros in the non-diagonal elements, it might not be strictly needed to occupy all that memory.

One alternative would be to return a sparse matrix that represents the identity matrix:

import numpy as np
from scipy.sparse import eye_array

import inversion_ideas as ii


class ModelNorm(ii.base.Objective):
    """
    Simple model norm regularization.

    Parameters
    ----------
    n_params : int
        Number of elements expected in the ``model`` array.
    """

    def __init__(self, n_params):
        self._n_params = n_params
        self._name = "norm"

    @property
    def n_params(self):
        return self._n_params

    def __call__(self, model):
        return model.T @ model

    def gradient(self, model):
        return 2 * model

    def hessian(self, model):
        return 2 * eye_array(self.n_params)
model_norm = ModelNorm(n_params)
model_norm.hessian(model)
<DIAgonal sparse array of dtype 'float64'
	with 5 stored elements (1 diagonals) and shape (5, 5)>

This strategy avoids storing all those non-diagonal zero elements in the matrix, but we are still storing all the twos in the diagonal. We could reduce the memory requirements even more by making the hessian method to return a LinearOperator objects that implements dot products in the same way as the identity matrix scaled by 2:

import numpy as np
from scipy.sparse.linalg import LinearOperator

import inversion_ideas as ii


class ModelNorm(ii.base.Objective):
    """
    Simple model norm regularization.

    Parameters
    ----------
    n_params : int
        Number of elements expected in the ``model`` array.
    """

    def __init__(self, n_params):
        self._n_params = n_params
        self._name = "norm"

    @property
    def n_params(self):
        return self._n_params

    def __call__(self, model):
        return model.T @ model

    def gradient(self, model):
        return 2 * model

    def hessian(self, model):

        def matvec(v):
            "Dot product of 2I with vector."
            return 2 * v

        def rmatvec(v):
            "Dot product of transpose of 2I with vector."
            return 2 * v

        hessian = LinearOperator(
            shape=(self.n_params, self.n_params),
            dtype=np.float64,
            matvec=matvec,
            rmatvec=rmatvec,
        )
        return hessian
model_norm = ModelNorm(n_params)
hessian = model_norm.hessian(model)
hessian
<5x5 _CustomLinearOperator with dtype=float64>

We can use this LinearOperator to compute dot products:

vector = np.random.default_rng(seed=52).random(n_params)

hessian @ vector
array([1.23152591, 1.14845259, 1.03908716, 0.01277908, 1.10058479])

Plugging custom objective functions#

Once we defined our objective function in a correct way, we can plug it into our inversion framework as with any other objective function from it.

For example, consider the same data_misfit term we defined in Objective function:

data_misfit
\[\phi_\text{d} (\mathbf{m})\]

Let’s build a ModelNorm regularization for this problem:

model_norm = ModelNorm(data_misfit.n_params)

We can build a linear combination for the two objective functions and choose a trade-off parameter for the model_norm:

phi = data_misfit + 30.0 * model_norm
phi
\[\phi_\text{d} (\mathbf{m}) + 30.0 \, \phi_\text{norm} (\mathbf{m})\]

We can minimize this objective function with the minimize.conjugate_gradient (since the problem is linear):

initial_model = np.zeros(data_misfit.n_params)

inverted_model = ii.minimize.conjugate_gradient(phi, initial_model)
inverted_model
array([-5.20251337,  1.72959343, -2.26550802,  4.04188586,  1.56209423])