Merge branch 'main' into model-wiring-matrix-slicer

Author: santisoler

src/inversion_ideas/base/objective_function.py

@@ -10,10 +10,9 @@
 
 import numpy as np
 import numpy.typing as npt
-from scipy.sparse import sparray
 from scipy.sparse.linalg import LinearOperator, aslinearoperator
 
-from ..typing import Model
+from ..typing import Model, SparseArray
 
 
 class Objective(ABC):
@@ -51,7 +50,7 @@ def gradient(self, model: Model) -> npt.NDArray[np.float64]:
     @abstractmethod
     def hessian(
         self, model: Model
-    ) -> npt.NDArray[np.float64] | sparray | LinearOperator:
+    ) -> npt.NDArray[np.float64] | SparseArray | LinearOperator:
         """
         Evaluate the hessian of the objective function for a given model.
         """
@@ -150,7 +149,7 @@ def gradient(self, model: Model) -> npt.NDArray[np.float64]:
 
     def hessian(
         self, model: Model
-    ) -> npt.NDArray[np.float64] | sparray | LinearOperator:
+    ) -> npt.NDArray[np.float64] | SparseArray | LinearOperator:
         """
         Evaluate the hessian of the objective function for a given model.
         """
@@ -243,7 +242,7 @@ def gradient(self, model: Model) -> npt.NDArray[np.float64]:
 
     def hessian(
         self, model: Model
-    ) -> npt.NDArray[np.float64] | sparray | LinearOperator:
+    ) -> npt.NDArray[np.float64] | SparseArray | LinearOperator:
         """
         Evaluate the hessian of the objective function for a given model.
         """
@@ -362,8 +361,8 @@ def _get_n_params(functions: list) -> int:
 
 
 def _sum(
-    operators: Iterator[npt.NDArray | sparray | LinearOperator],
-) -> npt.NDArray | sparray | LinearOperator:
+    operators: Iterator[npt.NDArray | SparseArray | LinearOperator],
+) -> npt.NDArray | SparseArray | LinearOperator:
     """
     Sum objects within an iterator.
 

src/inversion_ideas/data_misfit.py

@@ -4,13 +4,13 @@
 
 import numpy as np
 import numpy.typing as npt
-from scipy.sparse import dia_array, diags_array, sparray
+from scipy.sparse import dia_array, diags_array
 from scipy.sparse.linalg import LinearOperator, aslinearoperator
 
 from .base import Objective
-from .typing import Model
 from .utils import support_model_slice
 from .wires import ModelSlice
+from .typing import Model
 
 
 class DataMisfit(Objective):
@@ -112,7 +112,7 @@ def gradient(self, model: Model) -> npt.NDArray[np.float64]:
     @support_model_slice
     def hessian(
         self, model: Model
-    ) -> npt.NDArray[np.float64] | sparray | LinearOperator:
+    ) -> npt.NDArray[np.float64] | SparseArray | LinearOperator:
         """
         Hessian matrix.
         """

src/inversion_ideas/minimize/_functions.py

@@ -33,12 +33,12 @@ 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, sparray or LinearOperator or Callable, optional
+    preconditioner : (n_params, n_params) array, sparse array or LinearOperator or Callable, 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`.
+        array, sparse array or a `LinearOperator`.
     kwargs : dict
         Extra arguments that will be passed to the :func:`scipy.sparse.linalg.cg`
         function.

src/inversion_ideas/preconditioners.py

@@ -2,10 +2,10 @@
 Classes and functions to build preconditioners.
 """
 
-from scipy.sparse import diags_array, sparray
+from scipy.sparse import diags_array
 
 from .base import Objective
-from .typing import Model
+from .typing import Model, SparseArray
 
 
 class JacobiPreconditioner:
@@ -30,7 +30,7 @@ class JacobiPreconditioner:
     def __init__(self, objective_function: Objective):
         self.objective_function = objective_function
 
-    def __call__(self, model: Model) -> sparray:
+    def __call__(self, model: Model) -> SparseArray:
         """
         Generate a Jacobi preconditioner as a sparse diagonal array for a given model.
 

src/inversion_ideas/recipes.py

@@ -68,7 +68,7 @@ def create_l2_inversion(
         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 : {"jacobi"} or 2d array or sparse array 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``
@@ -184,7 +184,7 @@ def create_sparse_inversion(
         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 : {"jacobi"} or 2d array or sparse array 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``

src/inversion_ideas/typing.py

@@ -6,15 +6,43 @@
 
 import numpy as np
 import numpy.typing as npt
-from scipy.sparse import sparray
+from scipy.sparse import (
+    bsr_array,
+    bsr_matrix,
+    coo_array,
+    coo_matrix,
+    csc_array,
+    csc_matrix,
+    csr_array,
+    csr_matrix,
+    dia_array,
+    dia_matrix,
+)
 from scipy.sparse.linalg import LinearOperator
 
+SparseArray: TypeAlias = (
+    bsr_array
+    | bsr_matrix
+    | coo_array
+    | coo_matrix
+    | csc_array
+    | csc_matrix
+    | csr_array
+    | csr_matrix
+    | dia_array
+    | dia_matrix
+)
+"""
+Type alias to represent sparse arrays.
+"""
+
 Model: TypeAlias = npt.NDArray[np.float64]
 """
 Type alias to represent models in the inversion framework as 1D arrays.
 """
 
-Preconditioner: TypeAlias = npt.NDArray[np.float64] | sparray | LinearOperator
+
+Preconditioner: TypeAlias = npt.NDArray[np.float64] | SparseArray | LinearOperator
 """
 Type for static preconditioners.
 

src/inversion_ideas/utils.py

@@ -8,7 +8,8 @@
 
 import numpy as np
 import numpy.typing as npt
-from scipy.sparse import sparray
+
+from .typing import SparseArray
 
 from inversion_ideas.wires import ModelSlice
 
@@ -123,7 +124,7 @@ def wrapper(self, model, *args, **kwargs):
 def get_sensitivity_weights(
     jacobian: npt.NDArray[np.float64],
     *,
-    data_weights: npt.NDArray[np.float64] | sparray | None = None,
+    data_weights: npt.NDArray[np.float64] | SparseArray | None = None,
     volumes: npt.NDArray[np.float64] | None = None,
     vmin: float | None = 1e-12,
 ):

tests/test_objective_functions.py

@@ -1,6 +1,7 @@
 """
 Test operations for objective functions.
 """
+
 import pytest
 import numpy as np
 from scipy.sparse import diags_array, sparray