Add BFGS preconditioner

src/inversion_ideas/__init__.py

@@ -15,7 +15,11 @@
 from .inversion import Inversion
 from .inversion_log import InversionLog, InversionLogRich
 from .minimize import GaussNewtonConjugateGradient, conjugate_gradient
-from .preconditioners import JacobiPreconditioner, get_jacobi_preconditioner
+from .preconditioners import (
+    BFGSPreconditioner,
+    JacobiPreconditioner,
+    get_jacobi_preconditioner,
+)
 from .recipes import (
     create_l2_inversion,
     create_sparse_inversion,
@@ -25,6 +29,7 @@
 from .simulations import wrap_simulation
 
 __all__ = [
+    "BFGSPreconditioner",
     "ChiTarget",
     "ConvergenceWarning",
     "CustomCondition",

src/inversion_ideas/minimize/_functions.py

@@ -5,19 +5,18 @@
 """
 
 import warnings
-from collections.abc import Callable
 
 from scipy.sparse.linalg import cg
 
 from ..base import Objective
 from ..errors import ConvergenceWarning
-from ..typing import Model, Preconditioner
+from ..typing import CanBeUpdated, Model, Preconditioner
 
 
 def conjugate_gradient(
     objective: Objective,
     initial_model: Model,
-    preconditioner: Preconditioner | Callable[[Model], Preconditioner] | None = None,
+    preconditioner: Preconditioner | None = None,
     **kwargs,
 ) -> Model:
     r"""
@@ -33,12 +32,11 @@ def conjugate_gradient(
         Objective function to be minimized.
     initial_model : (n_params) array
         Initial model used to start the minimization.
-    preconditioner : (n_params, n_params) array, sparse array or LinearOperator or Callable, optional
+    preconditioner : (n_params, n_params) array, sparse array or LinearOperator, optional
         Matrix used as preconditioner in the conjugant gradient algorithm.
         If None, no preconditioner will be used.
-        A callable can be passed to build the preconditioner dynamically: such
-        callable should take a single ``initial_model`` argument and return an
-        array, sparse array or a `LinearOperator`.
+        If the preconditioner implements an ``initialize`` method, the preconditioner
+        will be initialized before being used in the conjugate gradient algorithm.
     kwargs : dict
         Extra arguments that will be passed to the :func:`scipy.sparse.linalg.cg`
         function.
@@ -65,8 +63,8 @@ def conjugate_gradient(
         raise ValueError(msg)
 
     if preconditioner is not None:
-        if callable(preconditioner):
-            preconditioner = preconditioner(initial_model)
+        if isinstance(preconditioner, CanBeUpdated):
+            preconditioner.update(initial_model)
         kwargs["M"] = preconditioner
 
     # TODO: maybe it would be nice to add a `is_linear` attribute to the objective

src/inversion_ideas/minimize/_minimizers.py

@@ -11,7 +11,7 @@
 
 from ..base import Condition, Minimizer, MinimizerResult, Objective
 from ..errors import ConvergenceWarning
-from ..typing import Model, Preconditioner
+from ..typing import CanBeUpdated, Model, Preconditioner
 from ..utils import CountCalls, Counter, get_logger
 from ._utils import backtracking_line_search
 
@@ -65,9 +65,7 @@ def __call__(
         objective: Objective,
         initial_model: Model,
         *,
-        preconditioner: (
-            Preconditioner | Callable[[Model], Preconditioner] | None
-        ) = None,
+        preconditioner: Preconditioner | None = None,
         callback: Callable[[MinimizerResult], None] | None = None,
     ) -> Generator[Model]:
         """
@@ -76,30 +74,23 @@ def __call__(
         Parameters
         ----------
         objective : Objective
-            Objective function that will get minimized.
+            Objective function to be minimized.
         initial_model : (n_params) array
-            Initial model to start the minimization.
-        preconditioner : (n_params, n_params) array, sparray or LinearOperator or Callable, optional
+            Initial model used to start the minimization.
+        preconditioner : (n_params, n_params) array, sparse array or LinearOperator, optional
             Matrix used as preconditioner in the conjugant gradient algorithm.
             If None, no preconditioner will be used.
-            A callable can be passed to build the preconditioner dynamically: such
-            callable should take a single ``initial_model`` argument and return an
-            array, `sparray` or a `LinearOperator`.
+            If the preconditioner implements an ``update`` method, the preconditioner
+            will be updated **before** every conjugate gradient minimization.
         callback : callable, optional
             Callable that gets called after each iteration.
         """
+        # Add the preconditioner to kwargs that will be passed to the cg function
         cg_kwargs = self.cg_kwargs.copy()
-
-        # Define a static preconditioner for all Gauss-Newton iterations
         if preconditioner is not None:
             if "M" in self.cg_kwargs:
                 msg = "Cannot simultanously pass `preconditioner` and `M`."
                 raise ValueError(msg)
-            preconditioner = (
-                preconditioner
-                if not callable(preconditioner)
-                else preconditioner(initial_model)
-            )
             cg_kwargs["M"] = preconditioner
 
         # Perform Gauss-Newton iterations
@@ -146,6 +137,13 @@ def __call__(
             if self.stopping_criterion is not None and self.stopping_criterion(model):
                 break
 
+            # Compute gradient and hessian
+            gradient, hessian = objective.gradient(model), objective.hessian(model)
+
+            # Update preconditioner before running cg
+            if isinstance(preconditioner, CanBeUpdated):
+                preconditioner.update(model)
+
             # Add a counter for conjugate-gradient iterations
             if (key := "callback") not in cg_kwargs:
                 counter = Counter()
@@ -160,6 +158,7 @@ def __call__(
 
             # Get number of cg iterations from callback
             cg_iters = cg_kwargs["callback"].counts
+
             # Raise warning if cg didn't converge
             cg_converged = cg_code == 0
             if not cg_converged:

src/inversion_ideas/operators.py

@@ -1,5 +1,5 @@
 """
-Custom LinearOperator classes and utilities.
+Custom :class:`scipy.sparse.linalg.LinearOperator` classes.
 """
 
 import numpy as np
@@ -9,6 +9,46 @@
 from inversion_ideas.typing import HasDiagonal, SparseArray
 
 
+class Identity(LinearOperator):
+    """
+    Identity matrix.
+
+    Represent the identity matrix of size ``n`` through
+    a :class:`scipy.sparse.linalg.LinearOperator`, to avoid storing a large diagonal
+    array full of ones.
+
+    Parameters
+    ----------
+    n : int
+        Size of the square matrix.
+    dtype : dtype, optional
+        Data type used for the linear operator.
+    """
+
+    def __init__(self, n: int, dtype=np.float64):
+        self._n = n
+        super().__init__(shape=(n, n), dtype=dtype)
+
+    @property
+    def n(self):
+        """
+        Size of the identity matrix.
+        """
+        return self._n
+
+    def _matvec(self, x):
+        return x.copy()
+
+    def _adjoint(self):
+        return self
+
+    def diagonal(self):
+        """
+        Diagonal of the matrix.
+        """
+        return np.ones(self.n, dtype=self.dtype)
+
+
 def get_diagonal(operator: npt.NDArray | SparseArray | LinearOperator):
     r"""
     Extract diagonal of a linear operator.