Improving scalability for large datasets

[VladimirShitov]

Hi! I'm using your package a lot for my PhD project, thanks for publishing it! I believe the most exciting discoveries can be made for large datasets. Unfortunately, the current setup of the package doesn't scale very well in terms of RAM usage. Especially this line of code makes me worried (and ruins my pipeline):

    def compute_distance_matrix_from_representations(
    ...
        pairwise_dists = np.zeros((n_cells, n_donors, n_donors))

For a dataset with ~1 Million cells and 1000 donors, it tries to allocate 8.75 TiB of memory. This makes MrVI barely applicable to the most exciting scenarios!

I wonder if users even need this tensor. Maybe it makes sense to aggregate distances somehow? For example, report one matrix of size n_donors * n_donors with an average distance. It could then be calculated in a significantly more memory-efficient way.

I'd be glad to hear your thoughts on this and see the improvement of your wonderful package.