Use the Model type alias on more places (#61)

Author: santisoler

src/inversion_ideas/base/conditions.py

@@ -12,6 +12,8 @@
 from rich.panel import Panel
 from rich.tree import Tree
 
+from ..typing import Model
+
 
 def _get_info_title(condition, model) -> str:
     """
@@ -30,9 +32,9 @@ class Condition(ABC):
     """
 
     @abstractmethod
-    def __call__(self, model) -> bool: ...
+    def __call__(self, model: Model) -> bool: ...
 
-    def update(self, model):  # noqa: B027
+    def update(self, model: Model):  # noqa: B027
         """
         Update the condition.
         """
@@ -48,7 +50,7 @@ def initialize(self):  # noqa: B027
         # necessary. The base class implements it to provide a common interface, even
         # for those children that don't implement it.
 
-    def info(self, model) -> Tree:
+    def info(self, model: Model) -> Tree:
         """
         Display information about the condition for a given model.
         """
@@ -86,17 +88,17 @@ def __init__(self, condition_a, condition_b):
         self.condition_b = condition_b
 
     @abstractmethod
-    def __call__(self, model) -> bool: ...
+    def __call__(self, model: Model) -> bool: ...
 
-    def update(self, model):
+    def update(self, model: Model):
         """
         Update the underlying conditions.
         """
         for condition in (self.condition_a, self.condition_b):
             if hasattr(condition, "update"):
                 condition.update(model)
 
-    def info(self, model) -> Tree:
+    def info(self, model: Model) -> Tree:
         status = self(model)
         checkbox = "x" if status else " "
         color = "green" if status else "red"
@@ -128,7 +130,7 @@ class LogicalAnd(_Mixin, Condition):
     Mixin condition for the AND operation between two other conditions.
     """
 
-    def __call__(self, model) -> bool:
+    def __call__(self, model: Model) -> bool:
         return self.condition_a(model) and self.condition_b(model)
 
 
@@ -137,7 +139,7 @@ class LogicalOr(_Mixin, Condition):
     Mixin condition for the OR operation between two other conditions.
     """
 
-    def __call__(self, model) -> bool:
+    def __call__(self, model: Model) -> bool:
         return self.condition_a(model) or self.condition_b(model)
 
 
@@ -146,5 +148,5 @@ class LogicalXor(_Mixin, Condition):
     Mixin condition for the XOR operation between two other conditions.
     """
 
-    def __call__(self, model) -> bool:
+    def __call__(self, model: Model) -> bool:
         return self.condition_a(model) ^ self.condition_b(model)

src/inversion_ideas/base/objective_function.py

@@ -5,6 +5,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Iterator
 from copy import copy
+from numbers import Real
 from typing import Self
 
 import numpy as np
@@ -56,7 +57,7 @@ def hessian(
         """
 
     @abstractmethod
-    def hessian_diagonal(self, model) -> npt.NDArray[np.float64]:
+    def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
         Diagonal of the Hessian.
         """
@@ -135,7 +136,7 @@ def n_params(self) -> int:
         """
         return self.function.n_params
 
-    def __call__(self, model):
+    def __call__(self, model: Model):
         """
         Evaluate the objective function.
         """
@@ -155,7 +156,7 @@ def hessian(
         """
         return self.multiplier * self.function.hessian(model)
 
-    def hessian_diagonal(self, model) -> npt.NDArray[np.float64]:
+    def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
         Diagonal of the Hessian.
         """
@@ -187,11 +188,11 @@ def _repr_latex_(self):
             phi_str = f"[ {phi_str} ]"
         return rf"${multiplier_str} \, {phi_str}$"
 
-    def __imul__(self, value) -> Self:
+    def __imul__(self, value: Real) -> Self:
         self.multiplier *= value
         return self
 
-    def __itruediv__(self, value) -> Self:
+    def __itruediv__(self, value: Real) -> Self:
         self.multiplier /= value
         return self
 
@@ -201,7 +202,7 @@ class Combo(Objective):
     Sum of objective functions.
     """
 
-    def __init__(self, functions):
+    def __init__(self, functions: list[Objective]):
         _get_n_params(functions)  # check if functions have the same n_params
         self._functions = functions
 
@@ -228,7 +229,7 @@ def n_params(self) -> int:
         """
         return _get_n_params(self.functions)
 
-    def __call__(self, model):
+    def __call__(self, model: Model):
         """
         Evaluate the objective function.
         """
@@ -248,7 +249,7 @@ def hessian(
         """
         return _sum(f.hessian(model) for f in self.functions)
 
-    def hessian_diagonal(self, model) -> npt.NDArray[np.float64]:
+    def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
         Diagonal of the Hessian.
         """

src/inversion_ideas/base/simulation.py

@@ -8,6 +8,8 @@
 from numpy.typing import NDArray
 from scipy.sparse.linalg import LinearOperator
 
+from ..typing import Model
+
 
 class Simulation(ABC):
     """
@@ -33,13 +35,13 @@ def n_data(self) -> int:
         """
 
     @abstractmethod
-    def __call__(self, model) -> NDArray[np.float64]:
+    def __call__(self, model: Model) -> NDArray[np.float64]:
         """
         Evaluate simulation for a given model.
         """
 
     @abstractmethod
-    def jacobian(self, model) -> NDArray[np.float64] | LinearOperator:
+    def jacobian(self, model: Model) -> NDArray[np.float64] | LinearOperator:
         """
         Jacobian matrix for a given model.
         """

src/inversion_ideas/conditions.py

@@ -30,7 +30,7 @@ class CustomCondition(Condition):
     def __init__(self, func: Callable[[Model], bool]):
         self.func = func
 
-    def __call__(self, model) -> bool:
+    def __call__(self, model: Model) -> bool:
         return self.func(model)
 
     @classmethod
@@ -87,14 +87,14 @@ def __init__(self, data_misfit: DataMisfit, chi_target=1.0):
         self.data_misfit = data_misfit
         self.chi_target = chi_target
 
-    def __call__(self, model) -> bool:
+    def __call__(self, model: Model) -> bool:
         """
         Check if condition has been met.
         """
         chi = self.data_misfit.chi_factor(model)
         return float(chi) < self.chi_target
 
-    def info(self, model) -> Tree:
+    def info(self, model: Model) -> Tree:
         tree = super().info(model)
         tree.add("Condition: chi < chi_target")
         tree.add(f"chi        = {self.data_misfit.chi_factor(model):.2e}")
@@ -148,20 +148,20 @@ def __init__(self, rtol: float = 1e-3, atol: float = 0.0):
         self.rtol = rtol
         self.atol = atol
 
-    def __call__(self, model) -> bool:
+    def __call__(self, model: Model) -> bool:
         if not hasattr(self, "previous"):
             return False
         diff = float(np.linalg.norm(model - self.previous))
         previous = float(np.linalg.norm(self.previous))
         return diff <= max(previous * self.rtol, self.atol)
 
-    def update(self, model):
+    def update(self, model: Model):
         """
         Cache model as the ``previous`` one.
         """
         self.previous = model
 
-    def info(self, model) -> Tree:
+    def info(self, model: Model) -> Tree:
         tree = super().info(model)
         diff = float(np.linalg.norm(model - self.previous))
         previous = float(np.linalg.norm(self.previous))
@@ -229,20 +229,20 @@ def __init__(self, objective_function: Objective, rtol: float = 1e-3, atol=0.0):
         self.rtol = rtol
         self.atol = atol
 
-    def __call__(self, model) -> bool:
+    def __call__(self, model: Model) -> bool:
         if not hasattr(self, "previous"):
             return False
         diff = abs(self.objective_function(model) - self.previous)
         previous = abs(self.previous)
         return diff <= max(previous * self.rtol, self.atol)
 
-    def update(self, model):
+    def update(self, model: Model):
         """
         Cache value of objective function with model as the ``previous`` one.
         """
         self.previous: float = float(self.objective_function(model))
 
-    def info(self, model) -> Tree:
+    def info(self, model: Model) -> Tree:
         tree = super().info(model)
         diff = abs(self.objective_function(model) - self.previous)
         previous = abs(self.previous)
@@ -253,7 +253,7 @@ def info(self, model) -> Tree:
         tree.add(f"atol               = {self.atol:.2e}")
         return tree
 
-    def ratio(self, model) -> float:
+    def ratio(self, model: Model) -> float:
         """
         Ratio ``|φ(m) - φ(m_prev)|/|φ(m_prev)|``.
         """

src/inversion_ideas/data_misfit.py

@@ -8,6 +8,7 @@
 from scipy.sparse.linalg import LinearOperator, aslinearoperator
 
 from .base import Objective
+from .typing import Model
 from .utils import cache_on_model
 
 
@@ -92,23 +93,25 @@ def __init__(
         self.set_name("d")
 
     @cache_on_model
-    def __call__(self, model) -> float:
+    def __call__(self, model: Model) -> float:
         # TODO:
         # Cache invalidation: we should clean the cache if data or uncertainties change.
         # Or they should be immutable.
         residual = self.residual(model)
         weights_matrix = self.weights_matrix
         return residual.T @ weights_matrix.T @ weights_matrix @ residual
 
-    def gradient(self, model) -> npt.NDArray[np.float64]:
+    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))
 
-    def hessian(self, model) -> npt.NDArray[np.float64] | sparray | LinearOperator:
+    def hessian(
+        self, model: Model
+    ) -> npt.NDArray[np.float64] | sparray | LinearOperator:
         """
         Hessian matrix.
         """
@@ -118,7 +121,7 @@ def hessian(self, model) -> npt.NDArray[np.float64] | sparray | LinearOperator:
         weights_matrix = aslinearoperator(self.weights_matrix)
         return 2 * jac.T @ weights_matrix.T @ weights_matrix @ jac
 
-    def hessian_diagonal(self, model) -> npt.NDArray[np.float64]:
+    def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
         Diagonal of the Hessian.
         """
@@ -146,7 +149,7 @@ def n_data(self):
         """
         return self.data.size
 
-    def residual(self, model):
+    def residual(self, model: Model):
         r"""
         Residual vector.
 
@@ -187,7 +190,7 @@ def weights_matrix(self) -> dia_array:
         """
         return diags_array(1 / self.uncertainty)
 
-    def chi_factor(self, model):
+    def chi_factor(self, model: Model):
         """
         Compute chi factor.
 

src/inversion_ideas/minimize/_minimizers.py

@@ -7,7 +7,6 @@
 from typing import Any
 
 import numpy as np
-import numpy.typing as npt
 from scipy.sparse.linalg import cg
 
 from ..base import Condition, Minimizer, Objective
@@ -66,7 +65,7 @@ def __call__(
         preconditioner: Preconditioner
         | Callable[[Model], Preconditioner]
         | None = None,
-    ) -> Generator[npt.NDArray[np.float64]]:
+    ) -> Generator[Model]:
         """
         Create iterator over Gauss-Newton minimization.
         """

src/inversion_ideas/regularization/_general.py

@@ -8,6 +8,7 @@
 
 from .._utils import prod_arrays
 from ..base import Objective
+from ..typing import Model
 
 
 class TikhonovZero(Objective):
@@ -62,7 +63,7 @@ def __init__(
         )
         self.set_name("0")
 
-    def __call__(self, model) -> float:
+    def __call__(self, model: Model) -> float:
         """
         Evaluate the regularization on a given model.
 
@@ -75,7 +76,7 @@ def __call__(self, model) -> float:
         weights_matrix = self.weights_matrix
         return model_diff.T @ weights_matrix.T @ weights_matrix @ model_diff
 
-    def gradient(self, model):
+    def gradient(self, model: Model):
         """
         Gradient vector.
 
@@ -88,7 +89,7 @@ def gradient(self, model):
         weights_matrix = self.weights_matrix
         return 2 * weights_matrix.T @ weights_matrix @ model_diff
 
-    def hessian(self, model):  # noqa: ARG002
+    def hessian(self, model: Model):  # noqa: ARG002
         """
         Hessian matrix.
 
@@ -100,7 +101,7 @@ def hessian(self, model):  # noqa: ARG002
         weights_matrix = self.weights_matrix
         return 2 * weights_matrix.T @ weights_matrix
 
-    def hessian_diagonal(self, model) -> npt.NDArray[np.float64]:
+    def hessian_diagonal(self, model: Model) -> npt.NDArray[np.float64]:
         """
         Diagonal of the Hessian.