Combinación de segmentación y indexación difundida para matrices numpy multidimensionales

Tengo una matriz de números ND (digamos, por ejemplo, 3x3x3) de la que me gustaría extraer una matriz secundaria, combinando segmentos y matrices de índices. Por ejemplo:

import numpy as np A = np.arange(3*3*3).reshape((3,3,3)) i0, i1, i2 = ([0,1], [0,1,2], [0,2]) ind1 = j0, j1, j2 = np.ix_(i0, i1, i2) ind2 = (j0, slice(None), j2) B1 = A[ind1] B2 = A[ind2] 

Yo esperaría que B1 == B2, pero en realidad, las formas son diferentes

 >>> B1.shape (2, 3, 2) >>> B2.shape (2, 1, 2, 3) >>> B1 array([[[ 0, 2], [ 3, 5], [ 6, 8]], [[ 9, 11], [12, 14], [15, 17]]]) >>> B2 array([[[[ 0, 3, 6], [ 2, 5, 8]]], [[[ 9, 12, 15], [11, 14, 17]]]]) 

Alguien entiende por qué? ¿Alguna idea de cómo podría obtener ‘B1’ manipulando solo los objetos ‘A’ e ‘ind2’? El objective es que funcione para cualquier nD matrices, y que no tenga que buscar la forma de las dimensiones que quiero mantener por completo (espero que sea lo suficientemente claro :)). ¡¡Gracias!!
—EDITAR—
Para ser más claros, me gustaría tener una función ‘divertida’ tal que

 A[fun(ind2)] == B1 

Esto es lo más cerca que puedo llegar a sus especificaciones, no he podido idear una solución que pueda calcular los índices correctos sin conocer A (o, más precisamente, su forma …).

 import numpy as np def index(A, s): ind = [] groups = s.split(';') for i, group in enumerate(groups): if group == ":": ind.append(range(A.shape[i])) else: ind.append([int(n) for n in group.split(',')]) return np.ix_(*ind) A = np.arange(3*3*3).reshape((3,3,3)) ind2 = index(A,"0,1;:;0,2") print A[ind2] 

Una version mas corta

 def index2(A,s):return np.ix_(*[range(A.shape[i])if g==":"else[int(n)for n in g.split(',')]for i,g in enumerate(s.split(';'))]) ind3 = index2(A,"0,1;:;0,2") print A[ind3] 

Los subespacios de indización de ind1 son (2,), (3,), (2,), y la B resultante es (2,3,2) . Este es un caso simple de indexación avanzada.

ind2 es un caso de indexación parcial (avanzada). Hay 2 matrices indexadas, y 1 rebanada. La documentación de indexación avanzada establece:

Si los subespacios de indexación están separados (por objetos de sector), entonces el espacio de indexación emitido es primero, seguido del subespacio segmentado de x.

En este caso, la indexación avanzada construye una matriz (2,2) (de los índices 1º y 3º) y agrega la dimensión de sector al final, lo que da como resultado una matriz (2,2,3) .

Explico el razonamiento con más detalle en https://stackoverflow.com/a/27097133/901925

Una forma de arreglar una tupla como ind2 , es expandir cada sector en una matriz. Recientemente vi esto hecho en np.insert .

 np.arange(*ind2[1].indices(3)) 

Se expande : hasta [0,1,2] . Pero el reemplazo tiene que tener la forma correcta.

 ind=list(ind2) ind[1]=np.arange(*ind2[1].indices(3)).reshape(1,-1,1) A[ind] 

Estoy dejando de lado los detalles para determinar qué término es una porción, su dimensión y la forma relevante. El objective es reproducir i1 .

Si los índices fueron generados por algo distinto de ix_ , la remodelación de esta porción podría ser más difícil. Por ejemplo

 A[np.array([0,1])[None,:,None],:,np.array([0,2])[None,None,:]] # (1,2,2,3) A[np.array([0,1])[None,:,None],np.array([0,1,2])[:,None,None],np.array([0,2])[None,None,:]] # (3,2,2) 

La división expandida debe ser compatible con las otras matrices en difusión.

Intercambiar ejes después de la indexación es otra opción. La lógica, sin embargo, podría ser más compleja. Pero en algunos casos la transposición podría ser más simple:

 A[np.array([0,1])[:,None],:,np.array([0,2])[None,:]].transpose(2,0,1) # (3,2,2) A[np.array([0,1])[:,None],:,np.array([0,2])[None,:]].transpose(0,2,1) # (2, 3, 2) 

En casos de indexación restringida como este con ix_ , es posible realizar la indexación en pasos sucesivos.

 A[ind1] 

es lo mismo que

 A[i1][:,i2][:,:,i3] 

y ya que i2 es el rango completo,

 A[i1][...,i3] 

Si solo tienes ind2 disponible.

 A[ind2[0].flatten()][[ind2[2].flatten()] 

En contextos más generales, debes saber cómo se emiten j0,j1,j2 entre sí, pero cuando son generados por ix_ , la relación es simple.

Puedo imaginar circunstancias en las que sería conveniente asignar A1 = A[i1] , seguido de una variedad de acciones que involucran a A1 , incluidas, entre otras, A1[...,i3] . Debe tener en cuenta cuándo A1 es una vista y cuándo es una copia.

Otra herramienta de indexación es take :

 A.take(i0,axis=0).take(i2,axis=2)