optim refactoring

* adding solver hook to get trainer, batch and losses befor optim
* moving self._instace in Optimizer

Author: dario-coscia

pina/optim/optimizer_interface.py

@@ -1,6 +1,7 @@
 """Module for the PINA Optimizer."""
 
 from abc import ABCMeta, abstractmethod
+from ..utils import check_consistency
 
 
 class Optimizer(metaclass=ABCMeta):
@@ -9,15 +10,56 @@ class Optimizer(metaclass=ABCMeta):
     should inherit form this class and implement the required methods.
     """
 
+    def __init__(self):
+        """
+        Initialization of the :class:`Optimizer` class.
+        """
+        self._extra_optim_args = {}
+        self._optimizer_instance = None
+
     @property
-    @abstractmethod
     def instance(self):
         """
-        Abstract property to retrieve the optimizer instance.
+        Get the optimizer instance.
+
+        :return: The optimizer instance.
+        :rtype: torch.optim.Optimizer
+        """
+        return self._optimizer_instance
+
+    @instance.setter
+    def instance(self, value):
+        """
+        Set the optimizer instance.
+
+        :param Any value: The optimizer instance.
         """
+        self._optimizer_instance = value
 
     @abstractmethod
     def hook(self):
         """
         Abstract method to define the hook logic for the optimizer.
         """
+
+    def get_optim_extra_args(self, trainer, batch, losses):
+        """
+        Retrieve and set extra optimizer arguments from the optimizer instance.
+
+        This method calls the ``get_optim_extra_args`` method of the underlying
+        optimizer instance (if it exists) and stores the resulting dictionary in
+        :attr:`extra_optim_args` of the optimizer instance.
+
+        :param trainer: The training manager controlling the optimization loop.
+        :type trainer: :class:`~pina.trainer.Trainer`
+        :param dict batch: The current batch of data used for training.
+        :param dict losses: Dictionary containing the computed loss values.
+        :raises ValueError: If ``extra_get_optim_extra_argsargs_dict`` does not
+            return a dictionary.
+        """
+        if hasattr(self.instance, "get_optim_extra_args"):
+            extra_args = self.instance.get_optim_extra_args(
+                trainer=trainer, batch=batch, losses=losses
+            )
+            check_consistency(extra_args, dict)
+            self.instance.extra_optim_args = extra_args

pina/optim/scheduler_interface.py

@@ -9,15 +9,31 @@ class Scheduler(metaclass=ABCMeta):
     inherit form this class and implement the required methods.
     """
 
+    def __init__(self):
+        """
+        Initialization of the :class:`Scheduler` class.
+        """
+        self._schedule_instance = None
+        self._last_lr = None
+
     @property
-    @abstractmethod
     def instance(self):
         """
         Abstract property to retrieve the scheduler instance.
         """
+        return self._schedule_instance
+
+    @instance.setter
+    def instance(self, value):
+        """
+        Set the optimizer instance.
+
+        :param Any value: The optimizer instance.
+        """
+        self._schedule_instance = value
 
     @abstractmethod
     def hook(self):
         """
         Abstract method to define the hook logic for the scheduler.
-        """
+        """

pina/optim/torch_optimizer.py

@@ -22,27 +22,13 @@ def __init__(self, optimizer_class, **kwargs):
             `here <https://pytorch.org/docs/stable/optim.html#algorithms>`_.
         """
         check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True)
-
         self.optimizer_class = optimizer_class
         self.kwargs = kwargs
-        self._optimizer_instance = None
 
     def hook(self, parameters):
         """
         Initialize the optimizer instance with the given parameters.
 
         :param dict parameters: The parameters of the model to be optimized.
         """
-        self._optimizer_instance = self.optimizer_class(
-            parameters, **self.kwargs
-        )
-
-    @property
-    def instance(self):
-        """
-        Get the optimizer instance.
-
-        :return: The optimizer instance.
-        :rtype: torch.optim.Optimizer
-        """
-        return self._optimizer_instance
+        self.instance = self.optimizer_class(parameters, **self.kwargs)

pina/optim/torch_scheduler.py

@@ -31,7 +31,6 @@ def __init__(self, scheduler_class, **kwargs):
 
         self.scheduler_class = scheduler_class
         self.kwargs = kwargs
-        self._scheduler_instance = None
 
     def hook(self, optimizer):
         """
@@ -40,16 +39,6 @@ def hook(self, optimizer):
         :param dict parameters: The parameters of the optimizer.
         """
         check_consistency(optimizer, Optimizer)
-        self._scheduler_instance = self.scheduler_class(
+        self.instance = self.scheduler_class(
             optimizer.instance, **self.kwargs
-        )
-
-    @property
-    def instance(self):
-        """
-        Get the scheduler instance.
-
-        :return: The scheduelr instance.
-        :rtype: torch.optim.LRScheduler
-        """
-        return self._scheduler_instance
+        )

pina/solver/solver.py

@@ -243,7 +243,12 @@ def _optimization_cycle(self, batch, **kwargs):
         :rtype: dict
         """
         # compute losses
-        losses = self.optimization_cycle(batch)
+        losses = self.optimization_cycle(batch, **kwargs)
+        # attach trainer, batch, losses to optimizer extra args
+        for optim in self._pina_optimizers:
+            optim.get_optim_extra_args(
+                trainer=self.trainer, batch=batch, losses=losses
+            )
         # clamp unknown parameters in InverseProblem (if needed)
         self._clamp_params()
         # store log