Implements the stationarity properties

Author: PierreQuinton

tests/unit/aggregation/_inputs.py

@@ -114,7 +114,7 @@ def _generate_weak_stationary_matrix(n_rows: int, n_cols: int) -> Tensor:
     _generate_strong_stationary_matrix(n_rows, n_cols)
     for n_rows, n_cols in _stationary_matrices_shapes
 ]
-weak_stationary_matrices = strong_stationary_matrices + [
+weak_stationary_matrices = [
     _generate_weak_stationary_matrix(n_rows, n_cols)
     for n_rows, n_cols in _stationary_matrices_shapes
 ]

tests/unit/aggregation/_property_testers.py

@@ -5,7 +5,13 @@
 
 from torchjd.aggregation import Aggregator
 
-from ._inputs import matrices, scaled_matrices, weak_stationary_matrices, zero_matrices
+from ._inputs import (
+    matrices,
+    scaled_matrices,
+    strong_stationary_matrices,
+    weak_stationary_matrices,
+    zero_matrices,
+)
 
 
 class ExpectedStructureProperty:
@@ -34,7 +40,7 @@ class NonConflictingProperty:
     """
 
     @classmethod
-    @mark.parametrize("matrix", weak_stationary_matrices + matrices)
+    @mark.parametrize("matrix", strong_stationary_matrices + weak_stationary_matrices + matrices)
     def test_non_conflicting_property(
         cls,
         aggregator: Aggregator,
@@ -80,3 +86,77 @@ def _assert_permutation_invariance_property(aggregator: Aggregator, matrix: Tens
     def _permute_randomly(matrix: Tensor) -> Tensor:
         row_permutation = torch.randperm(matrix.size(dim=0))
         return matrix[row_permutation]
+
+
+class StationarityProperty:
+    """
+    This class tests empirically that a given `Aggregator` satisfies the stationarity property.
+    """
+
+    @staticmethod
+    def _assert_stationarity_property(
+        aggregator: Aggregator,
+        stationary_matrix: Tensor,
+    ) -> None:
+        vector = aggregator(stationary_matrix)
+        norm = vector.norm().item()
+        assert norm < 8e-02
+
+    @staticmethod
+    def _assert_non_stationarity_property(
+        aggregator: Aggregator,
+        non_stationary_matrix: Tensor,
+    ) -> None:
+        vector = aggregator(non_stationary_matrix)
+        norm = vector.norm().item()
+        assert norm > 1e-03
+
+
+class StrongStationarityProperty(StationarityProperty):
+
+    @classmethod
+    @mark.parametrize("stationary_matrix", strong_stationary_matrices)
+    def test_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        stationary_matrix: Tensor,
+    ):
+        super(StrongStationarityProperty, cls)._assert_stationarity_property(
+            aggregator, stationary_matrix
+        )
+
+    @classmethod
+    @mark.parametrize("non_stationary_matrix", weak_stationary_matrices + matrices)
+    def test_non_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        non_stationary_matrix: Tensor,
+    ):
+        super(StrongStationarityProperty, cls)._assert_non_stationarity_property(
+            aggregator, non_stationary_matrix
+        )
+
+
+class WeakStationarityProperty(StationarityProperty):
+
+    @classmethod
+    @mark.parametrize("stationary_matrix", strong_stationary_matrices + weak_stationary_matrices)
+    def test_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        stationary_matrix: Tensor,
+    ):
+        super(WeakStationarityProperty, cls)._assert_stationarity_property(
+            aggregator, stationary_matrix
+        )
+
+    @classmethod
+    @mark.parametrize("non_stationary_matrix", matrices)
+    def test_non_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        non_stationary_matrix: Tensor,
+    ):
+        super(WeakStationarityProperty, cls)._assert_non_stationarity_property(
+            aggregator, non_stationary_matrix
+        )

tests/unit/aggregation/test_mgda.py

@@ -9,11 +9,17 @@
     ExpectedStructureProperty,
     NonConflictingProperty,
     PermutationInvarianceProperty,
+    WeakStationarityProperty,
 )
 
 
 @mark.parametrize("aggregator", [MGDA()])
-class TestMGDA(ExpectedStructureProperty, NonConflictingProperty, PermutationInvarianceProperty):
+class TestMGDA(
+    ExpectedStructureProperty,
+    NonConflictingProperty,
+    PermutationInvarianceProperty,
+    WeakStationarityProperty,
+):
     pass
 
 

tests/unit/aggregation/test_upgrad.py

@@ -7,11 +7,17 @@
     ExpectedStructureProperty,
     NonConflictingProperty,
     PermutationInvarianceProperty,
+    StrongStationarityProperty,
 )
 
 
 @mark.parametrize("aggregator", [UPGrad()])
-class TestUPGrad(ExpectedStructureProperty, NonConflictingProperty, PermutationInvarianceProperty):
+class TestUPGrad(
+    ExpectedStructureProperty,
+    NonConflictingProperty,
+    PermutationInvarianceProperty,
+    StrongStationarityProperty,
+):
     pass