Improve stationary matrices generation

Author: PierreQuinton

tests/unit/aggregation/_inputs.py

@@ -13,40 +13,53 @@ def _generate_matrix(m: int, n: int, rank: int) -> Tensor:
     return A
 
 
-def _generate_strong_stationary_matrix(m: int, n: int) -> Tensor:
+def _generate_strong_stationary_matrix(m: int, n: int, rank: int) -> Tensor:
     """
-    Generates a random matrix A of shape [m, n] with rank min(n, m - 1), such that there exists a
-    vector 0<v with v^T A = 0.
+    Generates a random matrix A of shape [m, n] with provided rank, such that there exists a vector
+    0<v with v^T A = 0.
     """
 
     v = torch.abs(torch.randn([m]))
-    return _generate_matrix_orthogonal_to_vector(v, n)
+    U1 = normalize(v, dim=0).unsqueeze(1)
+    U2 = _generate_semi_orthonormal_complement(U1)
+    U = torch.hstack([U1, U2])
+    Vt = _generate_orthonormal_matrix(n)
+    S = torch.diag(torch.abs(torch.randn([rank])))
+    A = U[:, 1 : rank + 1] @ S @ Vt[1 : rank + 1, :]
+    return A
 
 
-def _generate_weak_stationary_matrix(m: int, n: int) -> Tensor:
+def _generate_weak_non_strong_stationary_matrix(m: int, n: int, rank: int) -> Tensor:
     """
-    Generates a random matrix A of shape [m, n] with rank min(n, m - 1), such that there exists a
-    vector 0<=v with one coordinate equal to 0 and such that v^T A = 0.
-
-    Note that if multiple coordinates of v were equal to 0, the generated matrix would still be weak
-    stationary, but here we only set one of them to 0 for simplicity.
+    Generates a random matrix A of shape [m, n] with provided rank, such that there exists a vector
+    0<=v with one coordinate equal to 0 and such that v^T A = 0, and there is no vector 0<w with
+    w^T A = 0.
     """
 
     v = torch.abs(torch.randn([m]))
-    i = torch.randint(0, m, [])
-    v[i] = 0.0
-    return _generate_matrix_orthogonal_to_vector(v, n)
+    split_index = torch.randint(0, m, []).item()
+    shuffled_range = torch.randperm(m)
+    U1 = torch.zeros([m, 2])
+    U1[shuffled_range[:split_index], 0] = normalize(v[shuffled_range[:split_index]], dim=0)
+    U1[shuffled_range[split_index:], 1] = normalize(v[shuffled_range[split_index:]], dim=0)
+    U2 = _generate_semi_orthonormal_complement(U1)
+    U = torch.hstack([U1, U2])
+    Vt = _generate_orthonormal_matrix(n)
+    S = torch.diag(torch.abs(torch.randn([rank])))
+    A = U[:, 1 : rank + 1] @ S @ Vt[1 : rank + 1, :]
+    return A
 
 
-def _generate_matrix_orthogonal_to_vector(v: Tensor, n: int) -> Tensor:
+def _generate_non_weak_stationary_matrix(m: int, n: int, rank: int) -> Tensor:
     """
-    Generates a random matrix A of shape [len(v), n] with rank min(n, len(v) - 1) such that
-    v^T A = 0.
+    Generates a random matrix A of shape [m, n] with provided rank, such that there is no vector
+    0<=w with w^T A = 0.
     """
 
-    rank = min(n, len(v) - 1)
-    Q = normalize(v, dim=0).unsqueeze(1)
-    U = _generate_semi_orthonormal_complement(Q)
+    v = torch.abs(torch.randn([m]))
+    U1 = normalize(v, dim=0).unsqueeze(1)
+    U2 = _generate_semi_orthonormal_complement(U1)
+    U = torch.hstack([U1, U2])
     Vt = _generate_orthonormal_matrix(n)
     S = torch.diag(torch.abs(torch.randn([rank])))
     A = U[:, :rank] @ S @ Vt[:rank, :]
@@ -92,9 +105,11 @@ def _generate_semi_orthonormal_complement(Q: Tensor) -> Tensor:
     (9, 11),
 ]
 
-_stationary_matrices_shapes = [
-    (5, 3),
-    (9, 11),
+_stationary_matrices_triples = [
+    (5, 3, 3),
+    (9, 11, 6),
+    (7, 13, 2),
+    (3, 5, 3),
 ]
 
 _scales = [0.0, 1e-10, 1e3, 1e5, 1e10, 1e15]
@@ -106,12 +121,20 @@ def _generate_semi_orthonormal_complement(Q: Tensor) -> Tensor:
 scaled_matrices = [scale * matrix for scale in _scales for matrix in matrices]
 zero_matrices = [torch.zeros([m, n]) for m, n in _zero_matrices_shapes]
 strong_stationary_matrices = [
-    _generate_strong_stationary_matrix(m, n) for m, n in _stationary_matrices_shapes
+    _generate_strong_stationary_matrix(m, n, rank) for m, n, rank in _stationary_matrices_triples
 ]
-weak_stationary_matrices = [
-    _generate_weak_stationary_matrix(m, n) for m, n in _stationary_matrices_shapes
+weak_non_strong_stationary_matrices = [
+    _generate_weak_non_strong_stationary_matrix(m, n, rank)
+    for m, n, rank in _stationary_matrices_triples
 ]
-typical_matrices = zero_matrices + matrices + weak_stationary_matrices + strong_stationary_matrices
+non_weak_stationary_matrices = [
+    _generate_non_weak_stationary_matrix(m, n, rank) for m, n, rank in _stationary_matrices_triples
+]
+non_strong_stationary_matrices = weak_non_strong_stationary_matrices + non_weak_stationary_matrices
+typical_matrices = (
+    zero_matrices + matrices + strong_stationary_matrices + non_strong_stationary_matrices
+)
+
 
 scaled_matrices_2_plus_rows = [matrix for matrix in scaled_matrices if matrix.shape[0] >= 2]
 typical_matrices_2_plus_rows = [matrix for matrix in typical_matrices if matrix.shape[0] >= 2]