weighting refactory

Co-authored-by: Dario Coscia <dariocos99@gmail.com>

Author: GiovanniCanali

pina/loss/ntk_weighting.py

@@ -2,7 +2,7 @@
 
 import torch
 from .weighting_interface import WeightingInterface
-from ..utils import check_consistency
+from ..utils import check_consistency, in_range
 
 
 class NeuralTangentKernelWeighting(WeightingInterface):
@@ -20,32 +20,31 @@ class NeuralTangentKernelWeighting(WeightingInterface):
 
     """
 
-    def __init__(self, alpha=0.5):
+    def __init__(self, update_every_n_epochs=1, alpha=0.5):
         """
         Initialization of the :class:`NeuralTangentKernelWeighting` class.
 
+        :param int update_every_n_epochs: The number of training epochs between
+            weight updates. If set to 1, the weights are updated at every epoch.
+            Default is 1.
         :param float alpha: The alpha parameter.
         :raises ValueError: If ``alpha`` is not between 0 and 1 (inclusive).
         """
-        super().__init__()
+        super().__init__(update_every_n_epochs=update_every_n_epochs)
 
         # Check consistency
         check_consistency(alpha, float)
-        if alpha < 0 or alpha > 1:
-            raise ValueError("alpha should be a value between 0 and 1")
-
+        if not in_range(alpha, [0, 1], strict=False):
+            raise ValueError("alpha must be in range (0, 1).")
         # Initialize parameters
         self.alpha = alpha
-        self.weights = {}
-        self.default_value_weights = 1.0
+        self.history = {}
 
-    def aggregate(self, losses):
+    def weights_update(self, losses):
         """
-        Weight the losses according to the Neural Tangent Kernel algorithm.
+        Update the weights of the losses.
 
         :param dict(torch.Tensor) input: The dictionary of losses.
-        :return: The aggregation of the losses. It should be a scalar Tensor.
-        :rtype: torch.Tensor
         """
         # Define a dictionary to store the norms of the gradients
         losses_norm = {}
@@ -59,15 +58,10 @@ def aggregate(self, losses):
             losses_norm[condition] = grads.norm()
 
         # Update the weights
-        self.weights = {
-            condition: self.alpha
-            * self.weights.get(condition, self.default_value_weights)
+        return {
+            condition: self.alpha * self.history.get(condition, 1)
             + (1 - self.alpha)
             * losses_norm[condition]
             / sum(losses_norm.values())
             for condition in losses
         }
-
-        return sum(
-            self.weights[condition] * loss for condition, loss in losses.items()
-        )

pina/loss/scalar_weighting.py

@@ -4,22 +4,6 @@
 from ..utils import check_consistency
 
 
-class _NoWeighting(WeightingInterface):
-    """
-    Weighting scheme that does not apply any weighting to the losses.
-    """
-
-    def aggregate(self, losses):
-        """
-        Aggregate the losses.
-
-        :param dict losses: The dictionary of losses.
-        :return: The aggregated losses.
-        :rtype: torch.Tensor
-        """
-        return sum(losses.values())
-
-
 class ScalarWeighting(WeightingInterface):
     """
     Weighting scheme that assigns a scalar weight to each loss term.
@@ -36,28 +20,44 @@ def __init__(self, weights):
             dictionary, the default value is used.
         :type weights: float | int | dict
         """
-        super().__init__()
 
-        # Check consistency
+        super().__init__(update_every_n_epochs=1, aggregator="sum")
+
         check_consistency([weights], (float, dict, int))
 
-        # Weights initialization
-        if isinstance(weights, (float, int)):
+        if isinstance(weights, dict):
+            self.values = weights
+            self.default_value_weights = 1
+        elif isinstance(weights, (float, int)):
+            self.values = {}
             self.default_value_weights = weights
-            self.weights = {}
         else:
-            self.default_value_weights = 1.0
-            self.weights = weights
+            raise ValueError
 
-    def aggregate(self, losses):
+    def weights_update(self, losses):
         """
-        Aggregate the losses.
+        Update the weighting scheme based on the given losses.
+
+        This method must be implemented by subclasses. Its role is to update
+        the values in ``self.weights`` (a mapping from loss names to their
+        corresponding weights). The updated weights will then be used by
+        :meth:`aggregate` to compute the final aggregated loss.
 
-        :param dict losses: The dictionary of losses.
-        :return: The aggregated losses.
-        :rtype: torch.Tensor
+        :param losses: Dictionary mapping loss condition names to loss tensors.
+        :type losses: dict[str, torch.Tensor]
+        :return: Dictionary mapping loss names to their updated weight values.
+        :rtype: dict[str, float]
         """
-        return sum(
-            self.weights.get(condition, self.default_value_weights) * loss
-            for condition, loss in losses.items()
-        )
+        return {
+            condition: self.values.get(condition, self.default_value_weights)
+            for condition in losses.keys()
+        }
+
+
+class _NoWeighting(ScalarWeighting):
+    """
+    Weighting scheme that does not apply any weighting to the losses.
+    """
+
+    def __init__(self):
+        super().__init__(weights=1)

pina/loss/self_adaptive_weighting.py

@@ -2,7 +2,6 @@
 
 import torch
 from .weighting_interface import WeightingInterface
-from ..utils import check_positive_integer
 
 
 class SelfAdaptiveWeighting(WeightingInterface):
@@ -22,59 +21,35 @@ class SelfAdaptiveWeighting(WeightingInterface):
 
     """
 
-    def __init__(self, k=100):
+    def __init__(self, update_every_n_epochs=1):
         """
         Initialization of the :class:`SelfAdaptiveWeighting` class.
 
-        :param int k: The number of epochs after which the weights are updated.
-            Default is 100.
-
-        :raises ValueError: If ``k`` is not a positive integer.
+        :param int update_every_n_epochs: The number of training epochs between
+            weight updates. If set to 1, the weights are updated at every epoch.
+            Default is 1.
         """
-        super().__init__()
-
-        # Check consistency
-        check_positive_integer(value=k, strict=True)
+        super().__init__(update_every_n_epochs=update_every_n_epochs)
 
-        # Initialize parameters
-        self.k = k
-        self.weights = {}
-        self.default_value_weights = 1.0
-
-    def aggregate(self, losses):
+    def weights_update(self, losses):
         """
-        Weight the losses according to the self-adaptive algorithm.
+        Update the weights of the losses.
 
-        :param dict(torch.Tensor) losses: The dictionary of losses.
-        :return: The aggregation of the losses. It should be a scalar Tensor.
-        :rtype: torch.Tensor
+        :param dict(torch.Tensor) input: The dictionary of losses.
         """
-        # If weights have not been initialized, set them to 1
-        if not self.weights:
-            self.weights = {
-                condition: self.default_value_weights for condition in losses
-            }
-
-        # Update every k epochs
-        if self.solver.trainer.current_epoch % self.k == 0:
-
-            # Define a dictionary to store the norms of the gradients
-            losses_norm = {}
-
-            # Compute the gradient norms for each loss component
-            for condition, loss in losses.items():
-                loss.backward(retain_graph=True)
-                grads = torch.cat(
-                    [p.grad.flatten() for p in self.solver.model.parameters()]
-                )
-                losses_norm[condition] = grads.norm()
-
-            # Update the weights
-            self.weights = {
-                condition: sum(losses_norm.values()) / losses_norm[condition]
-                for condition in losses
-            }
-
-        return sum(
-            self.weights[condition] * loss for condition, loss in losses.items()
-        )
+        # Define a dictionary to store the norms of the gradients
+        losses_norm = {}
+
+        # Compute the gradient norms for each loss component
+        for condition, loss in losses.items():
+            loss.backward(retain_graph=True)
+            grads = torch.cat(
+                [p.grad.flatten() for p in self.solver.model.parameters()]
+            )
+            losses_norm[condition] = grads.norm()
+
+        # Update the weights
+        return {
+            condition: sum(losses_norm.values()) / losses_norm[condition]
+            for condition in losses
+        }

pina/loss/weighting_interface.py

@@ -1,6 +1,10 @@
 """Module for the Weighting Interface."""
 
 from abc import ABCMeta, abstractmethod
+from typing import final
+from ..utils import check_positive_integer, is_function
+
+_AGGREGATE_METHODS = {"sum": sum, "mean": lambda x: sum(x) / len(x)}
 
 
 class WeightingInterface(metaclass=ABCMeta):
@@ -9,19 +13,86 @@ class WeightingInterface(metaclass=ABCMeta):
     should inherit from this class.
     """
 
-    def __init__(self):
+    def __init__(
+        self,
+        update_every_n_epochs=1,
+        aggregator="sum",
+    ):
         """
         Initialization of the :class:`WeightingInterface` class.
+
+        :param int update_every_n_epochs: The number of training epochs between
+            weight updates. If set to 1, the weights are updated at every epoch.
+            This parameter is ignored by static weighting schemes. Default is 1.
+        :param aggregator: Aggregation method. Either:
+            - 'sum' → torch.sum
+            - 'mean' → torch.mean
+            - callable → custom aggregation function
         """
+        # Check consistency
+        check_positive_integer(value=update_every_n_epochs, strict=True)
+
+        if isinstance(aggregator, str):
+            if aggregator not in _AGGREGATE_METHODS:
+                raise ValueError(
+                    f"Invalid aggregator '{aggregator}'. Must be one of "
+                    f"{list(_AGGREGATE_METHODS.keys())}."
+                )
+            aggregator = _AGGREGATE_METHODS[aggregator]
+
+        elif not is_function(aggregator):
+            raise TypeError(
+                f"Aggregator must be either a string or a callable, "
+                f"got {type(aggregator).__name__}."
+            )
+
+        # Initialization
         self._solver = None
+        self.update_every_n_epochs = update_every_n_epochs
+        self.aggregator_fn = aggregator
+        self._weights = {}
 
     @abstractmethod
+    def weights_update(self, losses):
+        """
+        Update the weighting scheme based on the given losses.
+
+        This method must be implemented by subclasses. Its role is to update
+        the values in ``self.weights`` (a mapping from loss names to their
+        corresponding weights). The updated weights will then be used by
+        :meth:`aggregate` to compute the final aggregated loss.
+
+        :param losses: Dictionary mapping loss condition names to loss tensors.
+        :type losses: dict[str, torch.Tensor]
+        :return: Dictionary mapping loss names to their updated weight values.
+        :rtype: dict[str, float]
+        """
+
+    @final
     def aggregate(self, losses):
         """
-        Aggregate the losses.
+        Update the weights (if required) and aggregate the given losses.
+
+        This method first checks whether the weights should be updated based on
+        the current epoch and the ``update_every_n_epochs`` setting. If so, it
+        calls :meth:`weights_update` to refresh the weights. Afterwards, it
+        aggregates the (weighted) losses into a single scalar tensor using the
+        configured aggregator function. This method must not be override.
 
-        :param dict losses: The dictionary of losses.
+        :param losses: Dictionary mapping loss names to loss tensors.
+        :type losses: dict[str, torch.Tensor]
+        :return: The aggregated loss tensor.
+        :rtype: torch.Tensor
         """
+        # update weights
+        if self.solver.trainer.current_epoch % self.update_every_n_epochs == 0:
+            self._weights = self.weights_update(losses)
+
+        # aggregate
+        return self.aggregator_fn(
+            self._weights[condition] * loss
+            for condition, loss in losses.items()
+        )
 
     @property
     def solver(self):

pina/solver/garom.py

@@ -274,7 +274,7 @@ def validation_step(self, batch):
             condition_loss[condition_name] = self._loss_fn(
                 snapshots, snapshots_gen
             )
-        loss = self.weighting.aggregate(condition_loss)
+        loss = self.weighting.update_and_aggregate(condition_loss)
         self.store_log("val_loss", loss, self.get_batch_size(batch))
         return loss
 
@@ -297,7 +297,7 @@ def test_step(self, batch):
             condition_loss[condition_name] = self._loss_fn(
                 snapshots, snapshots_gen
             )
-        loss = self.weighting.aggregate(condition_loss)
+        loss = self.weighting.update_and_aggregate(condition_loss)
         self.store_log("test_loss", loss, self.get_batch_size(batch))
         return loss
 

pina/solver/physics_informed_solver/rba_pinn.py

@@ -266,7 +266,9 @@ def _optimization_cycle(self, batch, batch_idx, **kwargs):
         self._clamp_params()
 
         # aggregate
-        loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
+        loss = self.weighting.update_and_aggregate(losses).as_subclass(
+            torch.Tensor
+        )
 
         return loss
 

pina/solver/physics_informed_solver/self_adaptive_pinn.py

@@ -367,7 +367,9 @@ def _optimization_cycle(self, batch, batch_idx, **kwargs):
         self._clamp_params()
 
         # Aggregate
-        loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
+        loss = self.weighting.update_and_aggregate(losses).as_subclass(
+            torch.Tensor
+        )
 
         return loss