Fast brute force numpy array combination

Question

I want to extract the combinations of numpy arrays in this way:

from itertools import combinations
import numpy as np

X = np.array(np.random.randn(4, 6))
combination = np.array([X[:, [i, j]] for i, j in combinations(range(X.shape[1]), 2)])

Is there a way to speed this up? My array has a shape of 200x50000.

Answer 1

I decided to use an example of shape 20x5000. For which I could reduce the time from ~58 sec to ~3 sec i.e. almost a factor of 20.

def combinations(arr):
    n = arr.shape[0]
    a = np.broadcast_to(arr, (n, n))
    b = np.broadcast_to(arr.reshape(-1,1), (n, n))
    upper = np.tri(n, n, -1, dtype='bool').T
    return b[upper].reshape(-1), a[upper].reshape(-1)

X = np.array(np.random.randn(20,5000))
%timeit X[:, [*combinations(np.arange(X.shape[1]))]]
%timeit np.array([X[:, [i, j]] for i, j in itertools.combinations(range(X.shape[1]), 2)])

is giving me

3.2 s ± 29.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
57.8 s ± 2.35 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

And

combination = np.array([X[:, [i, j]] for i, j in itertools.combinations(range(X.shape[1]), 2)])
np.allclose(combination, np.moveaxis(X[:, [*combinations(np.arange(X.shape[1]))]], -1, 0))

confirms I am calculating the right thing. Just the axes are ordered differently.