`OrthogonalProcrustesAlignment` computes the full pairwise distance matrix before subsampling, leading to excessive memory usage

[onford]

Hi CEBRA team,

First, thank you for making CEBRA open source! While reproducing the multi-animal hippocampus example, I encountered a memory issue in OrthogonalProcrustesAlignment.

After inspecting the implementation, I noticed that the current workflow is:

1. compute the full pairwise label distance matrix: distance = self._distance(label, ref_label);
2. sort every row: target_idx = np.argsort(distance, axis=1)[:, :self.top_k]
3. construct X and Y
4. apply subsample (if provided)
5. solve the Orthogonal Procrustes problem

This means that subsample only reduces the size of the final Procrustes problem, but does not reduce the memory required for the pairwise distance computation.

For example, aligning datasets of sizes ref_label(10178, 3) and label(47431, 3) creates a distance matrix of shape (47431, 10178), which already occupies several GB of memory before subsampling is applied. On my machine(running code on WSL), the process is terminated by the Linux OOM killer before reaching the subsampling step.

From the documentation, I initially expected subsample to reduce the overall memory footprint, but in the current implementation it only affects the final Procrustes optimization.

Is this intended behavior?

If so, it might be worth mentioning in the documentation that subsample does not reduce the memory required for nearest-neighbor matching.

Thanks!