Chunk OrthogonalProcrustesAlignment top_k search to cut peak memory

cebra/data/helper.py

@@ -82,6 +82,15 @@ def _require_numpy_array(array: Union[npt.NDArray, torch.Tensor]):
     return array
 
 
+# Target block size for the ``top_k`` search: the rows of ``label`` are
+# processed in blocks, so only a ``(chunk_size, n_ref)`` distance matrix and
+# its argsort indices are held at once. 10 million elements is roughly 80 MB
+# per array in float64/intp. This is a target rather than a hard cap: a block
+# is always at least one row, so the working set never goes below
+# ``(1, n_ref)``.
+_PROCRUSTES_MAX_DISTANCE_ELEMENTS = 10_000_000
+
+
 class OrthogonalProcrustesAlignment:
     """Aligns two dataset by solving the orthogonal Procrustes problem.
 
@@ -219,10 +228,18 @@ def fit(
                 f"got ref_data:{data.shape[0]} samples and ref_labels:{label.shape[0]} samples."
             )
 
-        distance = self._distance(label, ref_label)
-
-        # keep indexes of the {self.top_k} labels the closest to the reference labels
-        target_idx = np.argsort(distance, axis=1)[:, :self.top_k]
+        # Search the top_k closest reference labels for each label, chunking
+        # over the rows of `label` so the full (n_label, n_ref) distance matrix
+        # is never materialized. Each row is scored against all references, so
+        # the per-row selection is identical to the unchunked computation.
+        n_label, n_ref = label.shape[0], ref_label.shape[0]
+        chunk_size = max(1, _PROCRUSTES_MAX_DISTANCE_ELEMENTS // n_ref)
+        target_idx = np.empty((n_label, self.top_k), dtype=np.intp)
+        for start in range(0, n_label, chunk_size):
+            stop = start + chunk_size
+            distance = self._distance(label[start:stop], ref_label)
+            target_idx[start:stop] = np.argsort(distance,
+                                                axis=1)[:, :self.top_k]
 
         if data.shape[1] != ref_data.shape[1]:
             raise ValueError(

tests/test_data_helper.py

@@ -359,3 +359,107 @@ def test_ensembling_performances(n_models=3):
     score = decoder.score(valid_embedding, valid_label)
 
     assert all(score_i < score for score_i in scores)
+
+
+@pytest.mark.parametrize("n_label, n_ref, dim, top_k", [
+    (50, 40, 4, 5),
+    (37, 60, 3, 1),
+    (100, 20, 2, 10),
+    (23, 7, 3, 7),
+])
+@pytest.mark.parametrize("label_dim", [1, 2])
+def test_orthogonal_alignment_chunking_matches_full(n_label, n_ref, dim, top_k,
+                                                    label_dim, monkeypatch):
+    """Chunked top_k search reproduces the full-matrix path (issue #307)."""
+    import scipy.linalg
+    rng = np.random.RandomState(0)
+    ref_data = rng.uniform(0, 1, (n_ref, dim))
+    data = rng.uniform(0, 1, (n_label, dim))
+    # distinct (non-tie) labels so the reference ordering is unambiguous
+    ref_label = (rng.permutation(n_ref).astype(float) +
+                 rng.uniform(0, 1e-3, n_ref))[:, None]
+    label = (rng.permutation(n_label).astype(float) +
+             rng.uniform(0, 1e-3, n_label))[:, None]
+    if label_dim == 1:
+        ref_label, label = ref_label.ravel(), label.ravel()
+
+    model = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k)
+    # independent reference: full distance matrix, then top_k per row
+    distance = model._distance(label.reshape(n_label, -1),
+                               ref_label.reshape(n_ref, -1))
+    full_idx = np.argsort(distance, axis=1)[:, :top_k]
+    X = data[:, None].repeat(top_k, axis=1).reshape(-1, dim)
+    Y = ref_data[full_idx].reshape(-1, dim)
+    ref_transform, _ = scipy.linalg.orthogonal_procrustes(X, Y)
+
+    # force many small chunks: 3 rows of `label` at a time
+    monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS",
+                        3 * n_ref)
+    chunked = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k)
+    chunked.fit(ref_data, data, ref_label, label)
+    np.testing.assert_allclose(chunked._transform,
+                               ref_transform,
+                               rtol=1e-10,
+                               atol=1e-10)
+
+    # force a single chunk: must be bit-identical to the chunked run
+    monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS",
+                        n_label * n_ref + 1)
+    single = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k)
+    single.fit(ref_data, data, ref_label, label)
+    np.testing.assert_array_equal(single._transform, chunked._transform)
+    np.testing.assert_array_equal(single.transform(data),
+                                  chunked.transform(data))
+
+
+def test_orthogonal_alignment_chunking_bit_identical_with_ties(monkeypatch):
+    """Row chunking is bit-identical, even with tied distances (#307)."""
+    rng = np.random.RandomState(1)
+    n_label, n_ref, dim, top_k = 60, 30, 3, 5
+    ref_data = rng.uniform(0, 1, (n_ref, dim))
+    data = rng.uniform(0, 1, (n_label, dim))
+    # integer labels with many duplicates -> tied distances
+    ref_label = rng.randint(0, 5, size=(n_ref, 1)).astype(float)
+    label = rng.randint(0, 5, size=(n_label, 1)).astype(float)
+
+    monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS",
+                        n_label * n_ref + 1)
+    single = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k)
+    single.fit(ref_data, data, ref_label, label)
+
+    monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS",
+                        4 * n_ref)
+    chunked = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k)
+    chunked.fit(ref_data, data, ref_label, label)
+    np.testing.assert_array_equal(single._transform, chunked._transform)
+
+
+def test_orthogonal_alignment_chunking_bounds_block_size(monkeypatch):
+    """The top_k search runs in bounded blocks, not one full matrix (#307)."""
+    rng = np.random.RandomState(2)
+    n_label, n_ref, dim, top_k = 50, 10, 3, 4
+    ref_data = rng.uniform(0, 1, (n_ref, dim))
+    data = rng.uniform(0, 1, (n_label, dim))
+    ref_label = rng.uniform(0, 1, (n_ref, 1))
+    label = rng.uniform(0, 1, (n_label, 1))
+
+    rows_per_call = []
+    original = cebra_data_helper.OrthogonalProcrustesAlignment._distance
+
+    def spy(self, label_i, label_j):
+        rows_per_call.append(label_i.shape[0])
+        return original(self, label_i, label_j)
+
+    # 3 rows of `label` per block
+    monkeypatch.setattr(cebra_data_helper, "_PROCRUSTES_MAX_DISTANCE_ELEMENTS",
+                        3 * n_ref)
+    monkeypatch.setattr(cebra_data_helper.OrthogonalProcrustesAlignment,
+                        "_distance", spy)
+
+    model = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k)
+    model.fit(ref_data, data, ref_label, label)
+
+    # the full (n_label, n_ref) distance matrix is never materialized
+    assert max(rows_per_call) <= 3
+    assert len(rows_per_call) == (n_label + 2) // 3
+    assert sum(rows_per_call) == n_label