N-D wersja itertools.combinations w numpy

Question

Chciałbym wprowadzić itertools.combinations dla numpy. Na podstawie this discussion Mam funkcję, która działa dla wejścia 1D:

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) 

a wyjście ma sens:

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]]) 

jednak, że najlepiej byłoby, gdybym mógł rozwinąć go do wejść ND, gdzie dodatkowe wymiary pozwalają tylko na szybkie wykonywanie wielu połączeń naraz. Więc, koncepcyjnie, jeśli combs([1, 2, 3], 2) wytwarza [1, 2], [1, 3], [2, 3] i wytwarza [4, 5], [4, 6], [5, 6], wówczas combs((1,2,3) and (4,5,6), 2) powinno wytworzyć [1, 2], [1, 3], [2, 3] and [4, 5], [4, 6], [5, 6], gdzie "i" po prostu reprezentuje równoległe wiersze lub kolumny (w zależności od tego, co ma sens). (A także do dodatkowych wymiarów)

nie jestem pewien:

1. Jak zrobić wymiary pracować w logiczny sposób, który jest zgodny z okazji pozostałe funkcje działają (jak jak niektóre funkcje NumPy mieć axis= Parametr i domyślna oś 0. Tak więc prawdopodobnie oś 0 powinna być tą, którą łączę wzdłuż, a wszystkie inne osie tylko reprezentują obliczenia równoległe?)
2. Jak uzyskać powyższy kod do pracy z ND (teraz otrzymuję ValueError: setting an array element with a sequence.)
3. Czy jest lepszy sposób na wykonanie dt = dtype([('', a.dtype)]*r)?

Answer 1

Nie wiem, jak to będzie działać na wydajność mądry, ale można zrobić kombinacje na tablicy indeksu, a następnie wyodrębnić rzeczywiste plastry tablic z np.take:

def combs_nd(a, r, axis=0): 
    a = np.asarray(a) 
    if axis < 0: 
     axis += a.ndim 
    indices = np.arange(a.shape[axis]) 
    dt = np.dtype([('', np.intp)]*r) 
    indices = np.fromiter(combinations(indices, r), dt) 
    indices = indices.view(np.intp).reshape(-1, r) 
    return np.take(a, indices, axis=axis) 

>>> combs_nd([1,2,3], 2) 
array([[1, 2], 
     [1, 3], 
     [2, 3]]) 
>>> combs_nd([[1,2,3],[4,5,6]], 2, axis=1) 
array([[[1, 2], 
     [1, 3], 
     [2, 3]], 

     [[4, 5], 
     [4, 6], 
     [5, 6]]]) 

Answer 2

Można użyć itertools.combinations() stworzyć tablica indeks, a następnie za pomocą ozdobnego indeksowanie numPy za:

import numpy as np 
from itertools import combinations, chain 
from scipy.misc 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] 

wyjściowy:

[[[ 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]]]