Add a log for the minimizers

src/inversion_ideas/base/__init__.py

@@ -4,7 +4,7 @@
 
 from .conditions import Condition
 from .directive import Directive
-from .minimizer import Minimizer
+from .minimizer import Minimizer, MinimizerResult
 from .objective_function import Combo, Objective, Scaled
 from .simulation import Simulation
 
@@ -13,6 +13,7 @@
     "Condition",
     "Directive",
     "Minimizer",
+    "MinimizerResult",
     "Objective",
     "Scaled",
     "Simulation",

src/inversion_ideas/base/minimizer.py

@@ -3,19 +3,50 @@
 """
 
 from abc import ABC, abstractmethod
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 
 from ..typing import Model
 from .objective_function import Objective
 
 
+class MinimizerResult(dict):
+    """
+    Dictionary to store results of a single minimization iteration.
+
+    This class is a child of ``dict``, but allows to access the values through
+    attributes.
+
+    Notes
+    -----
+    Inspired in the :class:`scipy.optimize.OptimizeResult`.
+    """
+
+    def __getattr__(self, name):
+        try:
+            return self[name]
+        except KeyError as e:
+            raise AttributeError(name) from e
+
+    __setattr__ = dict.__setitem__  # type: ignore[assignment]
+    __delattr__ = dict.__delitem__  # type: ignore[assignment]
+
+    def __dir__(self):
+        return list(self.keys())
+
+
 class Minimizer(ABC):
     """
     Base class to represent minimizers as generators.
     """
 
     @abstractmethod
-    def __call__(self, objective: Objective, initial_model: Model) -> Generator[Model]:
+    def __call__(
+        self,
+        objective: Objective,
+        initial_model: Model,
+        *,
+        callback: Callable[[MinimizerResult], None] | None = None,
+    ) -> Generator[Model]:
         """
         Minimize objective function.
 
@@ -25,6 +56,9 @@ def __call__(self, objective: Objective, initial_model: Model) -> Generator[Mode
             Objective function to be minimized.
         initial_model : (n_params) array
             Initial model used to start the minimization.
+        callback : callable, optional
+            Callable that gets called after each iteration.
+            Takes a :class:`inversion_ideas.base.MinimizerResult` as argument.
 
         Returns
         -------

src/inversion_ideas/inversion.py

@@ -7,11 +7,20 @@
 """
 
 import typing
+import warnings
 from collections.abc import Callable
 
+from rich.console import Group, RenderableType
+from rich.live import Live
+from rich.spinner import Spinner
+from rich.text import Text
+from rich.tree import Tree
+
+from inversion_ideas.errors import ConvergenceWarning
+
 from .base import Condition, Directive, Minimizer, Objective
-from .inversion_log import InversionLog, InversionLogRich
-from .typing import Model
+from .inversion_log import InversionLog, InversionLogRich, MinimizerLog
+from .typing import Log, Model
 from .utils import get_logger
 
 
@@ -40,11 +49,15 @@ class Inversion:
         no limit on the total amount of iterations.
     cache_models : bool, optional
         Whether to cache each model after each iteration.
-    log : InversionLog or bool, optional
-        Instance of :class:`InversionLog` to store information about the inversion.
+    log : Log or bool, optional
+        Instance of :class:`InversionLog` to store information about the inversion,
+        or any object that follows the :class:`inversion_ideas.typing.Log` protocol.
         If `True`, a default :class:`InversionLog` is going to be used.
         If `False`, no log will be assigned to the inversion, and :attr:`Inversion.log`
         will be ``None``.
+    log_minimizers : bool, optional
+        Whether to log the minimizers or not. Logging minimizers is only possible when
+        the ``minimizer`` is an instance of :class:`inversion_ideas.base.Minimizer``.
     minimizer_kwargs : dict, optional
         Extra arguments that will be passed to the ``minimizer`` when called.
     """
@@ -59,7 +72,8 @@ def __init__(
         stopping_criteria: Condition | Callable[[Model], bool],
         max_iterations: int | None = None,
         cache_models=False,
-        log: "InversionLog | bool" = True,
+        log: Log | InversionLog | bool = True,
+        log_minimizers: bool = True,
         minimizer_kwargs: dict | None = None,
     ):
         self.objective_function = objective_function
@@ -72,6 +86,7 @@ def __init__(
         if minimizer_kwargs is None:
             minimizer_kwargs = {}
         self.minimizer_kwargs = minimizer_kwargs
+        self._log_minimizers = log_minimizers
 
         # Assign log
         if log is False:
@@ -86,6 +101,11 @@ def __init__(
         # Assign model as a copy of the initial model
         self.model = initial_model.copy()
 
+        # TODO: Support for handling custom callbacks for the minimizer
+        if log is not None and "callback" in self.minimizer_kwargs:
+            msg = "Passing a custom callback for the minimizer is not yet supported."
+            raise NotImplementedError(msg)
+
     def __next__(self):
         """
         Run next iteration in the inversion.
@@ -108,17 +128,21 @@ def __next__(self):
 
         # Check for stopping criteria before trying to run the iteration
         if self.stopping_criteria(self.model):
-            get_logger().info(
+            get_logger().debug(
                 "🎉 Inversion successfully finished due to stopping criteria."
             )
+            self._stop_code = 0
             raise StopIteration
 
         # Check if maximum number of iterations have been reached
         if self.max_iterations is not None and self.counter >= self.max_iterations:
-            get_logger().info(
+            msg = (
                 "⚠️ Inversion finished after reaching maximum number of iterations "
                 f"({self.max_iterations})."
             )
+            get_logger().debug(msg)
+            warnings.warn(msg, ConvergenceWarning, stacklevel=2)
+            self._stop_code = 1
             raise StopIteration
 
         # Update stopping criteria (if necessary)
@@ -137,10 +161,18 @@ def __next__(self):
                 directive(self.model, self.counter)
 
         # Minimize objective function
+        # ---------------------------
         if isinstance(self.minimizer, Minimizer):
-            # Keep only the last model of the minimizer iterator
+            # Generate a new minimizer log for this iteration
+            minimizer_kwargs = self.minimizer_kwargs.copy()
+            if self.log is not None and self.log_minimizers:
+                minimizer_log = MinimizerLog()
+                self.minimizer_logs.append(minimizer_log)
+                minimizer_kwargs["callback"] = minimizer_log.update
+
+            # Unpack the generator and keep only the last model
             *_, model = self.minimizer(
-                self.objective_function, self.model, **self.minimizer_kwargs
+                self.objective_function, self.model, **minimizer_kwargs
             )
         else:
             model = self.minimizer(
@@ -161,8 +193,21 @@ def __next__(self):
         return self.model
 
     def __iter__(self):
+        self._stop_code = None
         return self
 
+    @property
+    def stop_code(self) -> int | None:
+        """
+        Code obtained after inversion stopped.
+
+        Code 0: the stopping criteria was met.
+        Code 1: the inversion stopped due to that the maximum number of iterations was
+        encountered..
+        Code ``None``: the inversion is still running or hasn't started yet.
+        """
+        return getattr(self, "_stop_code", None)
+
     @property
     def counter(self) -> int:
         """
@@ -185,6 +230,22 @@ def models(self) -> list:
             self._models = [self.initial_model]
         return self._models
 
+    @property
+    def log_minimizers(self) -> bool:
+        """Whether minimizers will be logged or not."""
+        return self._log_minimizers and isinstance(self.minimizer, Minimizer)
+
+    @property
+    def minimizer_logs(self) -> list[None | MinimizerLog] | None:
+        """
+        Logs of minimizers.
+        """
+        if not self.log_minimizers:
+            return None
+        if not hasattr(self, "_minimizer_logs"):
+            self._minimizer_logs = [None]
+        return self._minimizer_logs
+
     def run(self, show_log=True) -> Model:
         """
         Run the inversion.
@@ -195,11 +256,44 @@ def run(self, show_log=True) -> Model:
             Whether to show the ``log`` (if it's defined) during the inversion.
         """
         if show_log and self.log is not None:
-            if not hasattr(self.log, "live"):
+            if not isinstance(self.log, RenderableType):
+                # TODO: Add message
                 raise NotImplementedError()
-            with self.log.live() as live:
+
+            spinner = Spinner(
+                name="dots", text="Starting inversion...", style="green", speed=1
+            )
+            log = Tree(self.log) if self.log_minimizers else self.log
+            group = Group(log, spinner)
+
+            with Live(group, refresh_per_second=10) as live:
                 for _ in self:
-                    live.refresh()
+                    if self.log_minimizers:
+                        minimizer_log = self.minimizer_logs[self.counter]
+                        if minimizer_log is not None:
+                            renderable = minimizer_log.__rich__()
+                            renderable.title = (
+                                f"Minimizer log for iteration {self.counter}"
+                            )
+                            log.add(renderable)
+                    spinner.text = f"Running iteration {self.counter + 1}..."
+
+                # Finish inversion
+                group.renderables.pop(-1)  # remove spinner
+                if self.stop_code == 0:
+                    text = Text(
+                        "🎉 Inversion successfully finished due to stopping criteria."
+                    )
+                elif self.stop_code == 1:
+                    text = Text(
+                        "⚠️ Inversion finished after reaching maximum number of iterations "
+                        f"({self.max_iterations})."
+                    )
+                else:
+                    text = Text(f"⚠️ Invalid stop code '{self.stop_code}'.")
+                group.renderables.append(text)
+                live.refresh()
+
         else:
             for _ in self:
                 pass