python - pairwise - Versión ND de itertools.combinations en numpy

python combinations of two lists

Question

Me gustaría implementar itertools.combinations para numpy. Basado en esta discusión , tengo una función que funciona para la entrada 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)

y el resultado tiene sentido:

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

sin embargo, sería mejor si pudiera expandirlo a las entradas ND, donde las dimensiones adicionales simplemente le permiten hacer llamadas múltiples de forma rápida a la vez. Así que, conceptualmente, si los combs([1, 2, 3], 2) producen [1, 2], [1, 3], [2, 3] y combs([4, 5, 6], 2) producen [4, 5], [4, 6], [5, 6] , luego los combs((1,2,3) and (4,5,6), 2) deberían producir [1, 2], [1, 3], [2, 3] and [4, 5], [4, 6], [5, 6] donde "y" solo representa filas o columnas paralelas (lo que tenga sentido). (y del mismo modo para dimensiones adicionales)

No estoy seguro:

1. Cómo hacer que las dimensiones funcionen de una manera lógica que sea consistente con la forma en que otras funciones funcionan (como cómo algunas funciones numpy tienen un parámetro axis= y un valor predeterminado de axis 0. Entonces probablemente el eje 0 debería ser el que estoy combinando, y todos los demás ejes solo representan cálculos paralelos?)
2. Cómo hacer que el código anterior funcione con ND (ahora mismo obtengo ValueError: setting an array element with a sequence. ) ValueError: setting an array element with a sequence.
3. ¿Hay una mejor manera de hacer dt = dtype([('''', a.dtype)]*r) ?

Answer 1

No estoy seguro de cómo funcionará en cuanto a rendimiento, pero puede hacer las combinaciones en una matriz de índice, luego extraer las secciones reales de la matriz 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

Puede usar itertools.combinations() para crear la matriz de índices y luego usar la elegante indexación de 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]

salida:

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