Add implementation of sparse inversion (#28)

Add a new `SparseSmallness` regularization that makes use of Lawson's
equation to approximate an Lp norm with an L2 norm. Add a new `Irls`
directive that implements the IRLS algorithm (Fournier et al., 2019): it
updates the IRLS threshold in sparse regularizations and cools down the
beta multiplier based on the current value of the data misfit. Add a new
`create_sparse_inversion` constructor function. Add two notebooks that
run gravity sparse inversions.

Author: santisoler

notebooks/30_gravity-sparse-inversion.ipynb

@@ -5,15 +5,19 @@
 from . import base, typing, utils
 from ._version import __version__
 from .conditions import ChiTarget, CustomCondition, ModelChanged, ObjectiveChanged
-from .constructors import create_inversion
+from .constructors import create_inversion, create_sparse_inversion
 from .data_misfit import DataMisfit
-from .directives import MultiplierCooler, UpdateSensitivityWeights
+from .directives import (
+    Irls,
+    MultiplierCooler,
+    UpdateSensitivityWeights,
+)
 from .errors import ConvergenceWarning
 from .inversion import Inversion
 from .inversion_log import InversionLog, InversionLogRich
 from .minimize import GaussNewtonConjugateGradient, conjugate_gradient
 from .preconditioners import JacobiPreconditioner, get_jacobi_preconditioner
-from .regularization import Flatness, Smallness, TikhonovZero
+from .regularization import Flatness, Smallness, SparseSmallness, TikhonovZero
 from .simulations import wrap_simulation
 
 __all__ = [
@@ -26,17 +30,20 @@
     "Inversion",
     "InversionLog",
     "InversionLogRich",
+    "Irls",
     "JacobiPreconditioner",
     "ModelChanged",
     "MultiplierCooler",
     "ObjectiveChanged",
     "Smallness",
+    "SparseSmallness",
     "TikhonovZero",
     "UpdateSensitivityWeights",
     "__version__",
     "base",
     "conjugate_gradient",
     "create_inversion",
+    "create_sparse_inversion",
     "get_jacobi_preconditioner",
     "typing",
     "utils",

notebooks/31_gravity-sparse-inversion-constructor.ipynb

@@ -81,6 +81,9 @@ class ChiTarget(Condition):
     """
 
     def __init__(self, data_misfit: DataMisfit, chi_target=1.0):
+        if not hasattr(data_misfit, "chi_factor"):
+            msg = "Invalid `data_misfit`: missing `chi_factor` method."
+            raise TypeError(msg)
         self.data_misfit = data_misfit
         self.chi_target = chi_target
 
@@ -250,6 +253,18 @@ def info(self, model) -> Tree:
         tree.add(f"atol               = {self.atol:.2e}")
         return tree
 
+    def ratio(self, model) -> float:
+        """
+        Ratio ``|φ(m) - φ(m_prev)|/|φ(m_prev)|``.
+        """
+        if not hasattr(self, "previous"):
+            return np.nan
+        diff = abs(self.objective_function(model) - self.previous)
+        previous = abs(self.previous)
+        if previous == 0.0:
+            return np.inf
+        return diff / previous
+
     def initialize(self):
         """
         Initialize condition and clean ``previous`` attribute.

src/inversion_ideas/__init__.py

@@ -8,9 +8,11 @@
 import numpy.typing as npt
 
 from .base import Minimizer, Objective
-from .conditions import ChiTarget
-from .directives import MultiplierCooler
+from .conditions import ChiTarget, ObjectiveChanged
+from .data_misfit import DataMisfit
+from .directives import Irls, MultiplierCooler
 from .inversion import Inversion
+from .inversion_log import Column
 from .preconditioners import JacobiPreconditioner
 from .typing import Model, Preconditioner
 
@@ -43,7 +45,7 @@ def create_inversion(
         Model norm :math:`\phi_m`.
     starting_beta : float
         Starting value for the trade-off parameter :math:`\beta`.
-    initial_model : npt.NDArray[np.float64]
+    initial_model : (n_params) array
         Initial model to use in the inversion.
     minimizer : Minimizer
         Instance of :class:`Minimizer` used to minimize the objective function during
@@ -116,3 +118,144 @@ def create_inversion(
         minimizer_kwargs=minimizer_kwargs,
     )
     return inversion
+
+
+def create_sparse_inversion(
+    data_misfit: DataMisfit,
+    model_norm: Objective,
+    *,
+    starting_beta: float,
+    initial_model: Model,
+    minimizer: Minimizer | Callable[[Objective, Model], Model],
+    beta_cooling_factor: float = 2.0,
+    data_misfit_rtol=1e-1,
+    chi_l2_target: float = 1.0,
+    model_norm_rtol: float = 1e-3,
+    max_iterations: int | None = None,
+    cache_models: bool = True,
+    preconditioner: Preconditioner | Callable[[Model], Preconditioner] | None = None,
+) -> Inversion:
+    r"""
+    Create sparse norm inversion of the form: :math:`\phi_d + \beta \phi_m`.
+
+    Build an inversion where :math:`\phi_m` is a sparse norm regularization.
+    An IRLS algorithm will be applied, split in two stages.
+    The inversion will stop when the following inequality holds:
+
+    .. math::
+
+        \frac{|\phi_m^{(k)} - \phi_m^{(k-1)}|}{|\phi_m^{(k-1)}|} <  \eta_{\phi_m}
+
+    where :math:`\eta_{\phi_m}` is the ``model_norm_rtol``.
+
+    Parameters
+    ----------
+    data_misfit : Objective
+        Data misfit term :math:`\phi_d`.
+    model_norm : Objective
+        Model norm :math:`\phi_m`. It can be a single objective function term or a combo
+        containing multiple ones. At least one of them should be a sparse regularization
+        term.
+    starting_beta : float
+        Starting value for the trade-off parameter :math:`\beta`.
+    initial_model : (n_params) array
+        Initial model to use in the inversion.
+    minimizer : Minimizer
+        Instance of :class:`Minimizer` used to minimize the objective function during
+        the inversion.
+    beta_cooling_factor : float, optional
+        Cooling factor for the trade-off parameter :math:`\beta`. Every
+        ``beta_cooling_rate`` iterations, the :math:`\beta` will be _cooled down_ by
+        dividing it by the ``beta_cooling_factor``.
+    data_misfit_rtol : float, optional
+        Tolerance for the data misfit. This value is used to determine whether to cool
+        down the IRLS threshold or beta. See eq. 21 in Fournier and Oldenburg (2019).
+    chi_l2_target : float, optional
+        Chi factor target for the stage one (the L2 inversion). Once this chi target is
+        reached, the second stage starts.
+    model_norm_rtol : float, optional
+        Tolerance for the model norm. This value is used to determine if the inversion
+        should stop. See eq. 22 in Fournier and Oldenburg (2019).
+    max_iterations : int, optional
+        Max amount of iterations that will be performed. If ``None``, then there will be
+        no limit on the total amount of iterations.
+    cache_models : bool, optional
+        Whether to cache models after each iteration in the inversion.
+    preconditioner : {"jacobi"} or 2d array or sparray or LinearOperator or callable or None, optional
+        Preconditioner that will be passed to the ``minimizer`` on every call during the
+        inversion. The preconditioner can be a predefined 2d array, a sparse array or
+        a LinearOperator. Alternatively, it can be a callable that takes the ``model``
+        as argument and returns a preconditioner matrix (same types listed before). If
+        ``"jacobi"``, a default Jacobi preconditioner that will get updated on every
+        iteration will be defined for the inversion. If None, no preconditioner will be
+        passed.
+
+    Returns
+    -------
+    Inversion
+    """
+    # Define objective function
+    regularization = starting_beta * model_norm
+    objective_function = data_misfit + regularization
+
+    # Define IRLS directive
+    directives = [
+        Irls(
+            regularization,
+            data_misfit=data_misfit,
+            chi_l2_target=chi_l2_target,
+            beta_cooling_factor=beta_cooling_factor,
+            data_misfit_rtol=data_misfit_rtol,
+        )
+    ]
+
+    # Stopping criteria
+    smallness_not_changing = ObjectiveChanged(model_norm, rtol=model_norm_rtol)
+
+    # Preconditioner
+    minimizer_kwargs = {}
+    if preconditioner is not None:
+        if isinstance(preconditioner, str):
+            if preconditioner == "jacobi":
+                preconditioner = JacobiPreconditioner(objective_function)
+            else:
+                msg = f"Invalid preconditioner '{preconditioner}'."
+                raise ValueError(msg)
+        minimizer_kwargs["preconditioner"] = preconditioner
+
+    # Define inversion
+    inversion = Inversion(
+        objective_function,
+        initial_model,
+        minimizer,
+        directives=directives,
+        stopping_criteria=smallness_not_changing,
+        cache_models=cache_models,
+        max_iterations=max_iterations,
+        log=True,
+        minimizer_kwargs=minimizer_kwargs,
+    )
+
+    # Add extra columns to log
+    if inversion.log is not None:
+        # TODO: fix this in case that model norm is a combo
+        inversion.log.add_column(
+            "IRLS", lambda _, __: "active" if model_norm.irls else "inactive"
+        )
+        inversion.log.add_column(
+            "IRLS threshold",
+            Column(
+                title="ε",
+                callable=lambda _, __: model_norm.threshold,
+                fmt=None,
+            ),
+        )
+        inversion.log.add_column(
+            "model_norm_relative_diff",
+            Column(
+                title=r"|φm_(k) - φm_(k-1)|/|φm_(k-1)|",
+                callable=lambda _, model: smallness_not_changing.ratio(model),
+                fmt=None,
+            ),
+        )
+    return inversion