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

@@ -9,6 +9,8 @@
 
 from .base import Objective
 from .typing import Model, SparseArray
+from .utils import support_model_slice
+from .wires import ModelSlice, MultiSlice
 
 
 class DataMisfit(Objective):
@@ -78,44 +80,52 @@ def __init__(
         uncertainty: npt.NDArray[np.float64],
         simulation,
         *,
+        model_slice: ModelSlice | MultiSlice | None = None,
         build_hessian=False,
     ):
         # TODO: Check that the data and uncertainties have the size as ndata in the
         #       simulation.
         self.data = data
         self.uncertainty = uncertainty
         self.simulation = simulation
+        self.model_slice = model_slice
         self.build_hessian = build_hessian
         self.set_name("d")
 
+    @support_model_slice
     def __call__(self, model: Model) -> float:
-        residual = self.residual(model)
+        residual = self.simulation(model) - self.data
         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:
         """
         Hessian matrix.
         """
         jac = self.simulation.jacobian(model)
+        weights_matrix = aslinearoperator(self.weights_matrix)
         if not self.build_hessian:
             jac = aslinearoperator(jac)
-        weights_matrix = aslinearoperator(self.weights_matrix)
         return 2 * jac.T @ weights_matrix.T @ weights_matrix @ jac
 
+    @support_model_slice
     def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
-        Diagonal of the Hessian.
+        Approximated diagonal of the Hessian.
         """
         jac = self.simulation.jacobian(model)
         if isinstance(jac, LinearOperator):
@@ -132,6 +142,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
@@ -166,6 +178,8 @@ def residual(self, model: Model):
         where :math:`\mathbf{d}` is the vector with observed data, :math:`\mathcal{F}`
         is the forward model, and :math:`\mathbf{m}` is the model vector.
         """
+        if self.model_slice is not None:
+            model = self.model_slice.extract(model)
         return self.simulation(model) - self.data
 
     @property
@@ -176,7 +190,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/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

src/inversion_ideas/regularization/_mesh_based.py

@@ -11,6 +11,8 @@
 from .._utils import prod_arrays
 from ..base import Objective
 from ..typing import Model
+from ..utils import support_model_slice
+from ..wires import ModelSlice, MultiSlice
 
 
 class _MeshBasedRegularization(Objective):
@@ -24,6 +26,12 @@ class _MeshBasedRegularization(Objective):
 
     @property
     def n_params(self) -> int:
+        if (model_slice := getattr(self, "model_slice", None)) is not None:
+            return model_slice.full_size
+        return self.n_active
+
+    @property
+    def n_active(self) -> int:
         return int(np.sum(self.active_cells))
 
     @property
@@ -48,7 +56,7 @@ def cell_weights(
                 "It must be an array or a dictionary."
             )
             raise TypeError(msg)
-        if isinstance(value, np.ndarray) and value.size != self.n_params:
+        if isinstance(value, np.ndarray) and value.size != self.n_active:
             msg = (
                 f"Invalid cell_weights array with '{value.size}' elements. "
                 f"It must have '{self.n_params}' elements, "
@@ -57,7 +65,7 @@ def cell_weights(
             raise ValueError(msg)
         if isinstance(value, dict):
             for key, array in value.items():
-                if array.size != self.n_params:
+                if array.size != self.n_active:
                     msg = (
                         f"Invalid cell_weights array '{key}' with "
                         f"'{array.size}' elements. "
@@ -81,13 +89,14 @@ class Smallness(_MeshBasedRegularization):
     active_cells : (n_cells) array or None, optional
         Array full of bools that indicate the active cells in the mesh. It must have the
         same amount of elements as cells in the mesh.
-    cell_weights : (n_params) array or dict of (n_params) arrays or None, optional
+    cell_weights : (n_active) array or dict of (n_active) arrays or None, optional
         Array with cell weights.
         For multiple cell weights, pass a dictionary where keys are strings and values
         are the different weights arrays.
         If None, no cell weights are going to be used.
-    reference_model : (n_params) array or None, optional
-        Array with values for the reference model.
+    reference_model : (n_active) array or None, optional
+        Array with values for the reference model. It must have the same number of
+        elements as active cells in the mesh.
 
     Notes
     -----
@@ -138,28 +147,32 @@ def __init__(
         active_cells: npt.NDArray[np.bool] | None = None,
         cell_weights: npt.NDArray | dict[str, npt.NDArray] | None = None,
         reference_model: Model | None = None,
+        model_slice: ModelSlice | MultiSlice | None = None,
     ):
         self.mesh = mesh
         self.active_cells = (
             active_cells
             if active_cells is not None
             else np.ones(self.mesh.n_cells, dtype=bool)
         )
+        # assign model_slice after active_cells so n_active is correct during __init__
+        self.model_slice = model_slice
 
         # Assign the cell weights through the setter
         self.cell_weights = (
             cell_weights
             if cell_weights is not None
-            else np.ones(self.n_params, dtype=np.float64)
+            else np.ones(self.n_active, dtype=np.float64)
         )
 
         self.reference_model = (
             reference_model
             if reference_model is not None
-            else np.zeros(self.n_params, dtype=np.float64)
+            else np.zeros(self.n_active, dtype=np.float64)
         )
         self.set_name("s")
 
+    @support_model_slice
     def __call__(self, model: Model) -> float:
         """
         Evaluate the regularization on a given model.
@@ -181,6 +194,7 @@ def __call__(self, model: Model) -> float:
             @ model_diff
         )
 
+    @support_model_slice
     def gradient(self, model: Model):
         """
         Gradient vector.
@@ -193,15 +207,17 @@ def gradient(self, model: Model):
         model_diff = model - self.reference_model
         weights_matrix = self.weights_matrix
         cell_volumes_sqrt = self._volumes_sqrt_matrix
-        return (
+        gradient = (
             2
             * cell_volumes_sqrt.T
             @ weights_matrix.T
             @ weights_matrix
             @ cell_volumes_sqrt
             @ model_diff
         )
+        return gradient
 
+    @support_model_slice
     def hessian(self, model: Model):  # noqa: ARG002
         """
         Hessian matrix.
@@ -213,13 +229,14 @@ def hessian(self, model: Model):  # noqa: ARG002
         """
         weights_matrix = self.weights_matrix
         cell_volumes_sqrt = self._volumes_sqrt_matrix
-        return (
+        hessian = (
             2
             * cell_volumes_sqrt.T
             @ weights_matrix.T
             @ weights_matrix
             @ cell_volumes_sqrt
         )
+        return hessian
 
     def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
@@ -233,7 +250,7 @@ def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         return self.hessian(model).diagonal()
 
     @property
-    def weights_matrix(self) -> dia_array:
+    def weights_matrix(self) -> dia_array[np.float64]:
         """
         Diagonal matrix with the square root of regularization weights on cells.
         """
@@ -247,7 +264,7 @@ def weights_matrix(self) -> dia_array:
         return diags_array(np.sqrt(cell_weights))
 
     @property
-    def _volumes_sqrt_matrix(self) -> dia_array:
+    def _volumes_sqrt_matrix(self) -> dia_array[np.float64]:
         """
         Diagonal matrix with the square root of cell volumes.
         """
@@ -270,12 +287,12 @@ class Flatness(_MeshBasedRegularization):
     active_cells : (n_cells) array or None, optional
         Array full of bools that indicate the active cells in the mesh. It must have the
         same amount of elements as cells in the mesh.
-    cell_weights : (n_params) array or dict of (n_params) arrays or None, optional
+    cell_weights : (n_active) array or dict of (n_params) arrays or None, optional
         Array with cell weights.
         For multiple cell weights, pass a dictionary where keys are strings and values
         are the different weights arrays.
         If None, no cell weights are going to be used.
-    reference_model : (n_params) array or None, optional
+    reference_model : (n_active) array or None, optional
         Array with values for the reference model.
 
     Notes
@@ -337,6 +354,7 @@ def __init__(
         active_cells: npt.NDArray[np.bool] | None = None,
         cell_weights: npt.NDArray | dict[str, npt.NDArray] | None = None,
         reference_model: Model | None = None,
+        model_slice: ModelSlice | MultiSlice | None = None,
     ):
         self.mesh = mesh
         self.direction = direction
@@ -345,21 +363,24 @@ def __init__(
             if active_cells is not None
             else np.ones(self.mesh.n_cells, dtype=bool)
         )
+        # assign model_slice after active_cells so n_active is correct during __init__
+        self.model_slice = model_slice
 
         # Assign the cell weights through the setter
         self.cell_weights = (
             cell_weights
             if cell_weights is not None
-            else np.ones(self.n_params, dtype=np.float64)
+            else np.ones(self.n_active, dtype=np.float64)
         )
 
         self.reference_model = (
             reference_model
             if reference_model is not None
-            else np.zeros(self.n_params, dtype=np.float64)
+            else np.zeros(self.n_active, dtype=np.float64)
         )
         self.set_name(direction)
 
+    @support_model_slice
     def __call__(self, model: Model) -> float:
         """
         Evaluate the regularization on a given model.
@@ -384,6 +405,7 @@ def __call__(self, model: Model) -> float:
             @ model_diff
         )
 
+    @support_model_slice
     def gradient(self, model: Model):
         """
         Gradient vector.
@@ -397,17 +419,19 @@ def gradient(self, model: Model):
         weights_matrix = self.weights_matrix
         cell_volumes_sqrt = self._volumes_sqrt_matrix
         cell_gradient = self._cell_gradient
-        return (
+        gradient = (
             2
-            * cell_gradient.T
+            @ cell_gradient.T
             @ cell_volumes_sqrt.T
             @ weights_matrix.T
             @ weights_matrix
             @ cell_volumes_sqrt
             @ cell_gradient
             @ model_diff
         )
+        return gradient
 
+    @support_model_slice
     def hessian(self, model: Model):  # noqa: ARG002
         """
         Hessian matrix.
@@ -420,7 +444,7 @@ def hessian(self, model: Model):  # noqa: ARG002
         weights_matrix = self.weights_matrix
         cell_gradient = self._cell_gradient
         cell_volumes_sqrt = self._volumes_sqrt_matrix
-        return (
+        hessian = (
             2
             * cell_gradient.T
             @ cell_volumes_sqrt.T
@@ -429,6 +453,7 @@ def hessian(self, model: Model):  # noqa: ARG002
             @ cell_volumes_sqrt
             @ cell_gradient
         )
+        return hessian
 
     def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """