Add model_slice to objective funcs to handle composed models

src/inversion_ideas/__init__.py

@@ -23,6 +23,7 @@
 )
 from .regularization import Flatness, Smallness, SparseSmallness, TikhonovZero
 from .simulations import wrap_simulation
+from .wires import Wires
 
 __all__ = [
     "ChiTarget",
@@ -43,6 +44,7 @@
     "SparseSmallness",
     "TikhonovZero",
     "UpdateSensitivityWeights",
+    "Wires",
     "__version__",
     "base",
     "conjugate_gradient",

src/inversion_ideas/data_misfit.py

@@ -12,6 +12,8 @@
 from .base import Objective
 from .operators import get_diagonal
 from .typing import Model, SparseArray
+from .utils import support_model_slice
+from .wires import ModelSlice, MultiSlice
 
 
 class DataMisfit(Objective):
@@ -96,6 +98,7 @@ def __init__(
         uncertainty: npt.NDArray[np.float64],
         simulation,
         *,
+        model_slice: ModelSlice | MultiSlice | None = None,
         build_hessian=False,
         estimate_hessian_diagonal=False,
     ):
@@ -104,23 +107,29 @@ def __init__(
         self.data = data
         self.uncertainty = uncertainty
         self.simulation = simulation
+        self.model_slice = model_slice
         self.build_hessian = build_hessian
         self.estimate_hessian_diagonal = estimate_hessian_diagonal
         self.set_name("d")
 
+    @support_model_slice
     def __call__(self, model: Model) -> float:
         residual = self.residual(model)
         weights_matrix = self.weights_matrix
         return residual.T @ weights_matrix.T @ weights_matrix @ residual
 
+    @support_model_slice
     def gradient(self, model: Model) -> npt.NDArray[np.float64]:
         """
         Gradient vector.
         """
         jac = self.simulation.jacobian(model)
         weights_matrix = self.weights_matrix
-        return 2 * jac.T @ (weights_matrix.T @ weights_matrix @ self.residual(model))
+        residual = self.simulation(model) - self.data
+        gradient = 2 * jac.T @ (weights_matrix.T @ weights_matrix @ residual)
+        return gradient
 
+    @support_model_slice
     def hessian(
         self, model: Model
     ) -> npt.NDArray[np.float64] | SparseArray | LinearOperator:
@@ -143,6 +152,7 @@ def hessian(
         weights_matrix = aslinearoperator(self.weights_matrix)
         return 2 * jac.T @ weights_matrix.T @ weights_matrix @ jac
 
+    @support_model_slice
     def hessian_approx(self, model: Model) -> npt.NDArray[np.float64] | SparseArray:
         """
         Approximated version of the Hessian.
@@ -180,6 +190,7 @@ def hessian_approx(self, model: Model) -> npt.NDArray[np.float64] | SparseArray:
             return self.hessian(model)  # type: ignore[return-value]
         return diags_array(self.hessian_diagonal(model))
 
+    @support_model_slice
     def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
         Get the main diagonal of the Hessian.
@@ -242,6 +253,8 @@ def n_params(self):
         """
         Number of model parameters.
         """
+        if self.model_slice is not None:
+            return self.model_slice.full_size
         return self.simulation.n_params
 
     @property
@@ -251,6 +264,7 @@ def n_data(self):
         """
         return self.data.size
 
+    @support_model_slice(expand_return=False)
     def residual(self, model: Model):
         r"""
         Residual vector.
@@ -286,7 +300,7 @@ def weights(self) -> npt.NDArray[np.float64]:
         return 1 / self.uncertainty**2
 
     @property
-    def weights_matrix(self) -> dia_array:
+    def weights_matrix(self) -> dia_array[np.float64]:
         """
         Diagonal matrix with the square root of the regularization weights.
         """

src/inversion_ideas/operators.py

@@ -4,9 +4,10 @@
 
 import numpy as np
 import numpy.typing as npt
+from scipy.sparse import dia_array
 from scipy.sparse.linalg import LinearOperator
 
-from inversion_ideas.typing import HasDiagonal, SparseArray
+from .typing import HasDiagonal, SparseArray
 
 
 def get_diagonal(operator: npt.NDArray | SparseArray | LinearOperator):
@@ -101,3 +102,105 @@ def _get_standard_basis(ndim: int, dtype=np.float64):
         vector = np.zeros(ndim, dtype=dtype)
         vector[i] = 1
         yield vector
+
+
+class BlockSquareMatrix(LinearOperator):
+    r"""
+    Operator that represents a square matrix with a non-zero block surrounded by zeros.
+
+    Use this class to represent hessian matrices that are built from smaller matrices
+    that operate only on a subset of the entire model. Only a block of that hessian will
+    be non-zero (the one related to the relevant model elements for that objective
+    function), while the rest of the matrix will be filled of zeros.
+
+    Parameters
+    ----------
+    block : array, sparse array or LinearOperator
+        Square block matrix.
+    slice_matrix : dia_array
+        Diagonal array to represent the slicing matrix.
+
+    Notes
+    -----
+    This ``LinearOperator`` represents square matrices like:
+
+    .. math::
+
+        \mathbf{H} = \begin{bmatrix}
+            0 & 0          & 0\\
+            0 & \mathbf{B} & 0\\
+            0 & 0          & 0
+            \end{bmatrix}
+
+    where :math:`\mathbf{B}` is a non-zero square block matrix that sits in the diagonal
+    of :math:`\mathbf{H}`. The matrix :math:`\mathbf{H}` can be built as:
+
+    .. math::
+
+        \mathbf{H} = \mathbf{s}^T \mathbf{B} \mathbf{s}
+
+    where :math:`\mathbf{s}` is the *slicing matrix*.
+    """
+
+    def __init__(
+        self,
+        block: npt.NDArray | LinearOperator | SparseArray,
+        slice_matrix: dia_array,
+    ):
+        if block.shape[0] != block.shape[1]:
+            msg = (
+                f"Invalid block matrix '{block}' with shape '{block.shape}'. "
+                "It must be a square matrix."
+            )
+            raise ValueError(msg)
+
+        if slice_matrix.shape[0] != block.shape[1]:
+            msg = (
+                f"Invalid block '{block}' and slice_matrix '{slice_matrix}' with "
+                f"shapes '{block.shape}' and {slice_matrix.shape}."
+            )
+            raise ValueError(msg)
+
+        if slice_matrix.shape[1] <= block.shape[1]:
+            # TODO: improve msg
+            msg = "block is larger than slice matrix"
+            raise ValueError(msg)
+
+        _, full_size = slice_matrix.shape
+        shape = (full_size, full_size)
+        super().__init__(shape=shape, dtype=block.dtype)
+
+        self._block = block
+        self._slice_matrix = slice_matrix
+
+    @property
+    def block(self):
+        return self._block
+
+    @property
+    def slice_matrix(self) -> dia_array:
+        return self._slice_matrix
+
+    def _matvec(self, x):
+        """
+        Dot product between the matrix and a vector.
+        """
+        result = self.slice_matrix.T @ (self.block @ (self.slice_matrix @ x))
+        return result
+
+    def _rmatvec(self, x):
+        """
+        Dot product between the transposed matrix and a vector.
+        """
+        result = self.slice_matrix @ (self.block.T @ (self.slice_matrix.T @ x))
+        return result
+
+    def diagonal(self):
+        """
+        Diagonal of the matrix.
+        """
+        if not isinstance(self.block, HasDiagonal):
+            # TODO: Add msg
+            raise TypeError()
+        diagonal = self.block.diagonal()
+        return self.slice_matrix.T @ diagonal

src/inversion_ideas/recipes.py

@@ -16,6 +16,7 @@
 from .preconditioners import JacobiPreconditioner
 from .regularization import Flatness, Smallness
 from .typing import Model, Preconditioner
+from .wires import ModelSlice, MultiSlice
 
 
 def create_l2_inversion(
@@ -273,6 +274,7 @@ def create_tikhonov_regularization(
     alpha_y: float | None = None,
     alpha_z: float | None = None,
     reference_model_in_flatness: bool = False,
+    model_slice: ModelSlice | MultiSlice | None = None,
 ) -> Combo:
     """
     Create a linear combination of Tikhonov (L2) regularization terms.
@@ -321,13 +323,15 @@ def create_tikhonov_regularization(
         active_cells=active_cells,
         cell_weights=cell_weights,
         reference_model=reference_model,
+        model_slice=model_slice,
     )
     if alpha_s is not None:
         smallness = alpha_s * smallness
 
     kwargs = {
         "active_cells": active_cells,
         "cell_weights": cell_weights,
+        "model_slice": model_slice,
     }
     if reference_model_in_flatness:
         kwargs["reference_model"] = reference_model