feat(strategy): nystrom captures 3 subspaces, not 4 (drop redundant col(B) landmark block)

strategy/tests/test_nystrom.py

@@ -69,3 +69,51 @@ def test_nystrom_basis_flops_honest_and_cheaper_than_rsvd():
     assert ny == 2.0 * n * m * m
     assert ny > 0.0
     assert ny < RandomizedSVDTransform().basis_flops(n, m)
+
+
+def test_nystrom_recovers_low_rank_product():
+    # Regression for the #156-style fix: three-way split over
+    # {col(A), row(A), row(B)} must still recover an exact low-rank product,
+    # same shape as rsvd's post-#156 test_rsvd_recovers_low_rank_product.
+    backend = CPUBackend(verbose=False)
+    rng = np.random.default_rng(6)
+    n, r, m = 96, 5, 48                     # m/3 = 16 >= r per captured space
+    A = rng.standard_normal((n, r)) @ rng.standard_normal((r, n))
+    B = rng.standard_normal((n, r)) @ rng.standard_normal((r, n))
+    Q = backend.to_host(
+        NystromTransform(seed=0).basis(n, m, backend, np.float64, A=A, B=B)
+    )
+    C = Q @ (Q.T @ A @ Q) @ (Q.T @ B @ Q) @ Q.T
+    rel = np.linalg.norm(C - A @ B) / np.linalg.norm(A @ B)
+    assert rel < 1e-8
+
+
+def test_nystrom_col_b_is_redundant_but_col_a_and_row_b_are_not():
+    # Numeric proof using THIS transform's actual sketching method (landmark
+    # columns, not rsvd's random projections): col(A) and row(B) are
+    # load-bearing for Ĉ = P A P B P = A B, but col(B) never is. #156 proved
+    # this for random-projection sketches; this confirms the same algebra
+    # holds when each space is filled by landmark sampling instead, which is
+    # the actual justification for dropping col(B) here too.
+    rng = np.random.default_rng(9)
+    n, r = 80, 4
+    A = rng.standard_normal((n, r)) @ rng.standard_normal((r, n))
+    B = rng.standard_normal((n, r)) @ rng.standard_normal((r, n))
+    w = 5 * r  # generous width per space -- plenty to span an r-dim space
+
+    def landmark_block(X, seed, axis):
+        idx = np.random.default_rng(seed).choice(n, size=w, replace=False)
+        return X[:, idx] if axis == "col" else X[idx, :].T
+
+    col_a = landmark_block(A, 1, "col")
+    row_a = landmark_block(A, 2, "row")
+    row_b = landmark_block(B, 3, "row")
+
+    def rel_err(parts):
+        Q, _ = np.linalg.qr(np.concatenate(parts, axis=1))
+        C = Q @ (Q.T @ A @ Q) @ (Q.T @ B @ Q) @ Q.T
+        return np.linalg.norm(C - A @ B) / np.linalg.norm(A @ B)
+
+    assert rel_err([col_a, row_a, row_b]) < 1e-8   # all three spaces: exact
+    assert rel_err([row_a, row_b]) > 0.5            # missing col(A): breaks
+    assert rel_err([col_a, row_a]) > 0.5            # missing row(B): breaks

strategy/transforms.py

@@ -124,12 +124,19 @@ def basis_flops(self, n, m):
 class NystromTransform(Transform):
     """Landmark / Nyström column sampling over A and B.
 
-    Splits the M-column budget across col(A), row(A), col(B), and row(B) —
-    the same four spaces ``rsvd`` sketches — but forms each block by gathering
-    random landmark columns (or rows-as-columns) instead of random projections.
-    On genuine low-rank couples the landmarks span those spaces once enough
-    columns are drawn, so the thin QR that follows is enough; basis cost is
-    essentially the QR (``~2 N M²``), not ``rsvd``'s ``~2 N² M`` sketches.
+    Same three-space decomposition as ``rsvd`` (see ``RandomizedSVDTransform``'s
+    docstring for the algebra): the reconstruction ``Ĉ = P A P B P`` needs only
+    col(A), row(A), row(B) in range(Q) -- col(B) is provably redundant, and that
+    fact is about the projector, not about how each space gets sketched. This
+    used to split the M-column budget four ways (col(A), row(A), col(B), row(B))
+    like pre-#156 ``rsvd`` did, wasting a quarter of every landmark budget on a
+    space that never affects the result. Splits three ways now, over
+    {col(A), row(A), row(B)}, gathering random landmark columns (or
+    rows-as-columns) instead of random projections. On genuine low-rank couples
+    the landmarks span those spaces once enough columns are drawn, so the thin
+    QR that follows is enough; basis cost is essentially the QR (``~2 N M²``),
+    not ``rsvd``'s ``~2 N² M`` sketches. Total sketch width is still ``m``, so
+    ``basis_flops`` is unchanged by the three-way split.
     """
 
     name = "nystrom"
@@ -140,8 +147,8 @@ def basis(self, n, m, backend, dtype, A=None, B=None):
         if m < 1 or m > n:
             raise ValueError(f"nystrom requires 1 <= m <= n; got m={m}, n={n}")
 
-        base, rem = divmod(m, 4)
-        widths = [base + (1 if i < rem else 0) for i in range(4)]
+        base, rem = divmod(m, 3)
+        widths = [base + (1 if i < rem else 0) for i in range(3)]
         rng = np.random.default_rng(self.seed)
 
         def landmark_cols(X, w):
@@ -156,20 +163,19 @@ def landmark_rows_as_cols(X, w):
 
         parts = []
         if widths[0]:
-            parts.append(backend.to_device(landmark_cols(A, widths[0])))
+            parts.append(backend.to_device(landmark_cols(A, widths[0])))          # col(A)
         if widths[1]:
-            parts.append(backend.to_device(landmark_rows_as_cols(A, widths[1])))
+            parts.append(backend.to_device(landmark_rows_as_cols(A, widths[1])))  # row(A)
         if widths[2]:
-            parts.append(backend.to_device(landmark_cols(B, widths[2])))
-        if widths[3]:
-            parts.append(backend.to_device(landmark_rows_as_cols(B, widths[3])))
+            parts.append(backend.to_device(landmark_rows_as_cols(B, widths[2])))  # row(B)
 
         Y = backend.xp.concatenate(parts, axis=1)  # (n, m)
         return self._orthonormalize(Y, backend)
 
     def basis_flops(self, n, m):
         # Column/row gathers are memory traffic, not FLOPs. The mandatory cost
-        # is the thin QR of the (n, m) landmark stack (~2 n m²).
+        # is the thin QR of the (n, m) landmark stack (~2 n m²) -- unchanged by
+        # which three spaces filled the columns.
         return 2.0 * n * m * m