Improve input matrix generation

tests/unit/aggregation/_inputs.py

@@ -2,138 +2,76 @@
 from torch import Tensor
 
 
-def _check_valid_dimensions(n_rows: int, n_cols: int) -> None:
-    if n_rows < 1:
-        raise ValueError(
-            f"Parameter `n_rows` should be a positive integer. Found n_rows = {n_rows}."
-        )
-    if n_cols < 1:
-        raise ValueError(
-            f"Parameter `n_cols` should be a positive integer. Found n_cols = {n_cols}."
-        )
-
-
-def _check_valid_rank(n_rows: int, n_cols: int, rank: int) -> None:
-    if rank < 0:
-        raise ValueError(f"Parameter `rank` should be a non-negative integer. Found rank = {rank}.")
-    if rank > n_rows:
-        raise ValueError(
-            "Parameter `rank` should not be larger than the number of rows. "
-            f"Found rank = {rank} and n_rows = {n_rows}."
-        )
-    if rank > n_cols:
-        raise ValueError(
-            "Parameter `rank` should not be larger than the number of columns. "
-            f"Found rank = {rank} and n_cols = {n_cols}."
-        )
-
-
-def _augment_orthogonal_matrix(orthogonal_matrix: Tensor) -> Tensor:
-    """
-    Augments the provided matrix with one more column that is filled with a random unit vector that
-    is orthogonal to the provided orthogonal_matrix.
-    """
+def _generate_matrix(n_rows: int, n_cols: int, rank: int) -> Tensor:
+    """Generates a random matrix of shape [``n_rows``, ``n_cols``] with provided ``rank``."""
 
-    n_rows = orthogonal_matrix.shape[0]
-    projection = orthogonal_matrix @ orthogonal_matrix.T
-    zero = torch.zeros([n_rows])
-    while True:
-        random_vector = torch.randn([n_rows])
-        projected_vector = random_vector - projection @ random_vector
-        if not torch.allclose(projected_vector, zero):
-            break
-    projected_vector = torch.nn.functional.normalize(projected_vector, dim=0).reshape([-1, 1])
-    augmented_matrix = torch.cat((orthogonal_matrix, projected_vector), dim=1)
-    return augmented_matrix
+    U = _generate_orthogonal_matrix(n_rows)
+    Vt = _generate_orthogonal_matrix(n_cols)
+    S = torch.diag(torch.abs(torch.randn([rank])))
+    matrix = U[:, :rank] @ S @ Vt[:rank, :]
+    return matrix
 
 
-def _complete_orthogonal_matrix(orthogonal_matrix: Tensor, n_cols: int) -> Tensor:
+def _generate_strong_stationary_matrix(n_rows: int, n_cols: int) -> Tensor:
     """
-    Iteratively augments the input ``orthogonal_matrix`` with columns that are orthogonal to its
-    existing columns, until it has the required number of columns. Returns the obtained
-    orthogonal matrix.
+    Generates a random matrix of shape [``n_rows``, ``n_cols``] with rank
+    ``min(n_cols, n_rows - 1)``, such that there exists a vector ``0<v`` with ``v @ matrix = 0``.
     """
 
-    if orthogonal_matrix.shape[1] > n_cols:
-        raise ValueError(
-            f"Parameter `n_cols` should exceed the second dimension of the provided matrix. Found "
-            f"`n_cols = {n_cols}` and `partial_matrix.shape[1] = {orthogonal_matrix.shape[1]}`."
-        )
-
-    for i in range(n_cols - 1):
-        orthogonal_matrix = _augment_orthogonal_matrix(orthogonal_matrix)
-    return orthogonal_matrix
-
+    v = torch.abs(torch.randn([n_rows]))
+    return _generate_matrix_with_orthogonal_vector(v, n_cols)
 
-def _generate_unitary_matrix(n_rows: int, n_cols: int) -> Tensor:
-    """Generates a unitary matrix of shape [n_rows, n_cols]."""
 
-    _check_valid_dimensions(n_rows, n_cols)
-    partial_matrix = torch.randn([n_rows, 1])
-    partial_matrix = torch.nn.functional.normalize(partial_matrix, dim=0)
-
-    unitary_matrix = _complete_orthogonal_matrix(partial_matrix, n_cols)
-    return unitary_matrix
-
-
-def _generate_unitary_matrix_with_positive_column(n_rows: int, n_cols: int) -> Tensor:
+def _generate_weak_stationary_matrix(n_rows: int, n_cols: int) -> Tensor:
     """
-    Generates a unitary matrix of shape [n_rows, n_cols] with the first column consisting of an all
-    positive vector.
+    Generates a random matrix of shape [``n_rows``, ``n_cols``] with rank
+    ``min(n_cols, n_rows - 1)``, such that there exists a vector ``0<=v`` with at least one
+    coordinate equal to ``0`` and such that ``v @ matrix = 0``.
     """
-    _check_valid_dimensions(n_rows, n_cols)
-    partial_matrix = torch.abs(torch.randn([n_rows, 1]))
-    partial_matrix = torch.nn.functional.normalize(partial_matrix, dim=0)
 
-    unitary_matrix_with_positive_column = _complete_orthogonal_matrix(partial_matrix, n_cols)
-    return unitary_matrix_with_positive_column
+    v = torch.abs(torch.randn([n_rows]))
+    v[torch.randint(0, n_rows, [])] = 0.0
+    return _generate_matrix_with_orthogonal_vector(v, n_cols)
 
 
-def _generate_diagonal_singular_values(rank: int) -> Tensor:
-    """
-    generates a diagonal matrix of positive values sorted in descending order.
-    """
-    singular_values = torch.abs(torch.randn([rank]))
-    singular_values = torch.sort(singular_values, descending=True)[0]
-    S = torch.diag(singular_values)
-    return S
+def _generate_orthogonal_matrix(dim: int) -> Tensor:
+    """Uniformly generates a random orthogonal matrix of shape [``dim``, ``dim``]."""
 
+    A = torch.randn([dim, dim])
+    Q, _ = torch.linalg.qr(A)
+    return Q
 
-def generate_matrix(n_rows: int, n_cols: int, rank: int) -> Tensor:
+
+def _generate_matrix_with_orthogonal_vector(vector: Tensor, n_cols: int) -> Tensor:
     """
-    Generates a random matrix of shape [``n_rows``, ``n_cols``] with provided ``rank``.
+    Generates a random matrix of shape [``len(vector)``, ``n_cols``] with rank
+    ``min(n_cols, len(vector) - 1)`` such that ``vector @ matrix = 0``.
     """
 
-    _check_valid_rank(n_rows, n_cols, rank)
-
-    if rank == 0:
-        matrix = torch.zeros([n_rows, n_cols])
-    else:
-        U = _generate_unitary_matrix(n_rows, rank)
-        V = _generate_unitary_matrix(n_cols, rank)
-        S = _generate_diagonal_singular_values(rank)
-        matrix = U @ S @ V.T
-
+    n_rows = len(vector)
+    rank = min(n_cols, n_rows - 1)
+    U = _generate_orthogonal_complement(vector, n_cols=rank)
+    Vt = _generate_orthogonal_matrix(n_cols)
+    S = torch.diag(torch.abs(torch.randn([rank])))
+    matrix = U @ S @ Vt[:rank, :]
     return matrix
 
 
-def generate_stationary_matrix(n_rows: int, n_cols: int, rank: int) -> Tensor:
+def _generate_orthogonal_complement(vector: Tensor, n_cols: int) -> Tensor:
     """
-    Generates a random matrix of shape [``n_rows``, ``n_cols``] with provided ``rank``. The matrix
-    has a singular triple (u, s, v) such that u is all (strictly) positive and s is 0.
+    Uniformly generates a random orthogonal matrix of shape [``len(vector)``, ``n_cols``] such that
+    ``vector @ matrix = 0``.
     """
 
-    _check_valid_rank(n_rows, n_cols, rank)
-    if rank == 0:
-        matrix = torch.zeros([n_rows, n_cols])
-    else:
-        U = _generate_unitary_matrix_with_positive_column(n_rows, rank)
-        V = _generate_unitary_matrix(n_cols, rank)
-        S = _generate_diagonal_singular_values(rank)
-        S[0, 0] = 0.0
-        matrix = U @ S @ V.T
+    m = vector.shape[0]
+    u = torch.nn.functional.normalize(vector, dim=0)
+    A = torch.randn([m, n_cols])
 
-    return matrix
+    # project A onto the orthogonal complement of u
+    A_proj = A - u.unsqueeze(1) * (u.unsqueeze(0) @ A)
+
+    Q, _ = torch.linalg.qr(A_proj)
+    return Q
 
 
 _matrix_dimension_triples = [
@@ -145,27 +83,36 @@ def generate_stationary_matrix(n_rows: int, n_cols: int, rank: int) -> Tensor:
     (9, 11, 9),
 ]
 
-_zero_rank_matrix_shapes = [
+_zero_matrices_shapes = [
     (1, 1),
     (4, 3),
     (9, 11),
 ]
 
+_stationary_matrices_shapes = [
+    (5, 3),
+    (9, 11),
+]
+
 _scales = [0.0, 1e-10, 1.0, 1e3, 1e5, 1e10, 1e15]
 
 # Fix seed to fix randomness of matrix generation
 torch.manual_seed(0)
 
 matrices = [
-    generate_matrix(n_rows, n_cols, rank) for n_rows, n_cols, rank in _matrix_dimension_triples
+    _generate_matrix(n_rows, n_cols, rank) for n_rows, n_cols, rank in _matrix_dimension_triples
 ]
 scaled_matrices = [scale * matrix for scale in _scales for matrix in matrices]
-zero_rank_matrices = [torch.zeros([n_rows, n_cols]) for n_rows, n_cols in _zero_rank_matrix_shapes]
-matrices_2_plus_rows = [matrix for matrix in matrices + zero_rank_matrices if matrix.shape[0] >= 2]
+zero_matrices = [torch.zeros([n_rows, n_cols]) for n_rows, n_cols in _zero_matrices_shapes]
+matrices_2_plus_rows = [matrix for matrix in matrices + zero_matrices if matrix.shape[0] >= 2]
 scaled_matrices_2_plus_rows = [
-    matrix for matrix in scaled_matrices + zero_rank_matrices if matrix.shape[0] >= 2
+    matrix for matrix in scaled_matrices + zero_matrices if matrix.shape[0] >= 2
+]
+strong_stationary_matrices = [
+    _generate_strong_stationary_matrix(n_rows, n_cols)
+    for n_rows, n_cols in _stationary_matrices_shapes
 ]
-stationary_matrices = [
-    generate_stationary_matrix(n_rows, n_cols, rank)
-    for n_rows, n_cols, rank in _matrix_dimension_triples
+weak_stationary_matrices = strong_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,7 @@
 
 from torchjd.aggregation import Aggregator
 
-from ._inputs import matrices, scaled_matrices, stationary_matrices, zero_rank_matrices
+from ._inputs import matrices, scaled_matrices, weak_stationary_matrices, zero_matrices
 
 
 class ExpectedStructureProperty:
@@ -17,7 +17,7 @@ class ExpectedStructureProperty:
     """
 
     @classmethod
-    @mark.parametrize("matrix", scaled_matrices + zero_rank_matrices)
+    @mark.parametrize("matrix", scaled_matrices + zero_matrices)
     def test_expected_structure_property(cls, aggregator: Aggregator, matrix: Tensor):
         cls._assert_expected_structure_property(aggregator, matrix)
 
@@ -34,7 +34,7 @@ class NonConflictingProperty:
     """
 
     @classmethod
-    @mark.parametrize("matrix", stationary_matrices + matrices)
+    @mark.parametrize("matrix", weak_stationary_matrices + matrices)
     def test_non_conflicting_property(
         cls,
         aggregator: Aggregator,

tests/unit/aggregation/test_cagrad.py

@@ -4,7 +4,7 @@
 
 from torchjd.aggregation import CAGrad, Mean
 
-from ._inputs import matrices, stationary_matrices
+from ._inputs import matrices, strong_stationary_matrices
 from ._property_testers import ExpectedStructureProperty, NonConflictingProperty
 
 
@@ -20,7 +20,7 @@ class TestCAGradNonConflicting(NonConflictingProperty):
     pass
 
 
-@mark.parametrize("matrix", stationary_matrices + matrices)
+@mark.parametrize("matrix", strong_stationary_matrices + matrices)
 def test_equivalence_mean(matrix: Tensor):
     """Tests that CAGrad is equivalent to Mean when c=0."""
 

tests/unit/aggregation/test_constant.py

@@ -4,7 +4,7 @@
 
 from torchjd.aggregation import Constant
 
-from ._inputs import matrices, scaled_matrices, stationary_matrices, zero_rank_matrices
+from ._inputs import matrices, scaled_matrices, strong_stationary_matrices, zero_matrices
 from ._property_testers import ExpectedStructureProperty
 
 # The weights must be a vector of length equal to the number of rows in the matrix that it will be
@@ -19,10 +19,10 @@ def _make_aggregator(matrix: Tensor) -> Constant:
     return Constant(weights)
 
 
-_matrices_1 = scaled_matrices + zero_rank_matrices
+_matrices_1 = scaled_matrices + zero_matrices
 _aggregators_1 = [_make_aggregator(matrix) for matrix in _matrices_1]
 
-_matrices_2 = matrices + stationary_matrices
+_matrices_2 = matrices + strong_stationary_matrices
 _aggregators_2 = [_make_aggregator(matrix) for matrix in _matrices_2]