Versione N-D di itertools.combinazioni in numpy

Question

Vorrei implementare itertools.combinations per numpy. Sulla base di this discussion, ho una funzione che funziona per ingresso 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) 

e l'uscita senso:

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

tuttavia, sarebbe meglio se ho potuto espandersi a ingressi ND, dove le dimensioni aggiuntive ti consentono semplicemente di eseguire più chiamate contemporaneamente in modo rapido. Quindi, concettualmente, se combs([1, 2, 3], 2) produce [1, 2], [1, 3], [2, 3] e combs([4, 5, 6], 2) produce [4, 5], [4, 6], [5, 6], quindi combs((1,2,3) and (4,5,6), 2) dovrebbe produrre [1, 2], [1, 3], [2, 3] and [4, 5], [4, 6], [5, 6] dove "e" rappresenta solo righe o colonne parallele (a seconda del caso). (E allo stesso modo per ulteriori dimensioni)

non sono sicuro:

1. Come rendere le dimensioni funzionano in modo logico che è coerente con il modo in cui le altre funzioni funzionano (come come alcune funzioni NumPy hanno un axis= parametro e un valore predefinito dell'asse 0. Quindi probabilmente l'asse 0 dovrebbe essere quello che sto combinando e tutti gli altri assi rappresentano solo calcoli paralleli?)
2. Come ottenere il codice sopra per lavorare con ND (in questo momento ricevo ValueError: setting an array element with a sequence.)
3. C'è un modo migliore per fare dt = dtype([('', a.dtype)]*r)?

Answer 1

Non so come andrà a lavorare fuori prestazioni-saggio, ma si può fare le combinazioni su un array di indice, quindi estrarre le fette di matrice reale con 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

È possibile utilizzare itertools.combinations() per creare l'array di indici, e quindi utilizzare fantasia indicizzazione NumPy:

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] 

uscita:

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