N-D version of itertools.combinations in numpy

Question

I would like to implement itertools.combinations for numpy. Based on this discussion, I have a function that works for 1D input:

def combs(a, r):     """     Return successive r-length combinations of elements in the array a.     Should produce the same output as array(list(combinations(a, r))), but      faster.     """     a = asarray(a)     dt = dtype([('', a.dtype)]*r)     b = fromiter(combinations(a, r), dt)     return b.view(a.dtype).reshape(-1, r) 

and the output makes sense:

In [1]: list(combinations([1,2,3], 2)) Out[1]: [(1, 2), (1, 3), (2, 3)]  In [2]: array(list(combinations([1,2,3], 2))) Out[2]:  array([[1, 2],        [1, 3],        [2, 3]])  In [3]: combs([1,2,3], 2) Out[3]:  array([[1, 2],        [1, 3],        [2, 3]]) 

However, it would be best if I could expand it to N-D inputs, where additional dimensions simply allow you to speedily do multiple calls at once. So, conceptually, if combs([1, 2, 3], 2) produces [1, 2], [1, 3], [2, 3], and combs([4, 5, 6], 2) produces [4, 5], [4, 6], [5, 6], then combs((1,2,3) and (4,5,6), 2) should produce [1, 2], [1, 3], [2, 3] and [4, 5], [4, 6], [5, 6] where "and" just represents parallel rows or columns (whichever makes sense). (and likewise for additional dimensions)

I'm not sure:

1. How to make the dimensions work in a logical way that's consistent with the way other functions work (like how some numpy functions have an axis= parameter, and a default of axis 0. So probably axis 0 should be the one I am combining along, and all other axes just represent parallel calculations?)
2. How to get the above code to work with ND (right now I get ValueError: setting an array element with a sequence.)
3. Is there a better way to do dt = dtype([('', a.dtype)]*r)?

Answer 1

You can use itertools.combinations() to create the index array, and then use NumPy's fancy indexing:

import numpy as np from itertools import combinations, chain from scipy.special import comb  def comb_index(n, k):     count = comb(n, k, exact=True)     index = np.fromiter(chain.from_iterable(combinations(range(n), k)),                          int, count=count*k)     return index.reshape(-1, k)  data = np.array([[1,2,3,4,5],[10,11,12,13,14]])  idx = comb_index(5, 3) print(data[:, idx]) 

output:

[[[ 1  2  3]   [ 1  2  4]   [ 1  2  5]   [ 1  3  4]   [ 1  3  5]   [ 1  4  5]   [ 2  3  4]   [ 2  3  5]   [ 2  4  5]   [ 3  4  5]]   [[10 11 12]   [10 11 13]   [10 11 14]   [10 12 13]   [10 12 14]   [10 13 14]   [11 12 13]   [11 12 14]   [11 13 14]   [12 13 14]]]