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,139 @@ def test_ensembling_performances(n_models=3):
     score = decoder.score(valid_embedding, valid_label)
 
     assert all(score_i < score for score_i in scores)
+
+
+def _original_orthogonal_procrustes_fit(model, ref_data, data, ref_label,
+                                        label):
+    """Pre-#309 ``OrthogonalProcrustesAlignment.fit``, copied verbatim.
+
+    Every line below is copied from the implementation in commit
+    ``d1842cc``, with ``self`` renamed to ``model`` and the transform
+    returned instead of stored. The input validation and ``subsample``
+    branches of the original are omitted: #309 leaves them untouched and
+    the tests here never exercise them (labels always given,
+    ``subsample=None``).
+    """
+    import scipy.linalg
+
+    if len(ref_label.shape) == 1:
+        ref_label = np.expand_dims(ref_label, axis=1)
+    if len(label.shape) == 1:
+        label = np.expand_dims(label, axis=1)
+
+    distance = model._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)[:, :model.top_k]
+
+    # Get the whole data to align and only the selected closest samples
+    # from the reference data.
+    X = data[:, None].repeat(model.top_k, axis=1).reshape(-1, data.shape[1])
+    Y = ref_data[target_idx].reshape(-1, ref_data.shape[1])
+
+    # Compute orthogonal matrix that most closely maps X to Y using the orthogonal Procrustes problem.
+    transform, _ = scipy.linalg.orthogonal_procrustes(X, Y)
+
+    return transform
+
+
+@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),
+        # edge cases: single label row, single reference, tiny shapes
+        (1, 8, 3, 2),
+        (6, 1, 2, 1),
+        (2, 2, 2, 2),
+    ])
+@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 pre-#309 implementation."""
+    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()
+
+    # oracle: the verbatim pre-PR (full-matrix) implementation
+    oracle = cebra_data_helper.OrthogonalProcrustesAlignment(top_k=top_k)
+    ref_transform = _original_orthogonal_procrustes_fit(oracle, ref_data, data,
+                                                        ref_label, label)
+
+    # 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_array_equal(chunked._transform, ref_transform)
+
+    # force a single chunk: must also be bit-identical
+    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, ref_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