python - Indice n matriz dimensional con(n-1) matriz d

numpy

Question

¿Cuál es la forma más elegante de acceder a una matriz de n dimensiones con una matriz de dimensiones (n-1) a lo largo de una dimensión dada como en el ejemplo ficticio

a = np.random.random_sample((3,4,4))
b = np.random.random_sample((3,4,4))
idx = np.argmax(a, axis=0)

¿Cómo puedo acceder ahora con idx a para obtener los máximos en a como si hubiera usado a.max(axis=0) ? o cómo recuperar los valores especificados por idx en b ?

Pensé en usar np.meshgrid pero creo que es una exageración. Tenga en cuenta que el axis dimensión puede ser cualquier eje útil (0,1,2) y no se conoce de antemano. ¿Hay alguna forma elegante de hacer esto?

Answer 1

Utilice advanced-indexing :

m,n = a.shape[1:]
I,J = np.ogrid[:m,:n]
a_max_values = a[idx, I, J]
b_max_values = b[idx, I, J]

Para el caso general:

def argmax_to_max(arr, argmax, axis):
    """argmax_to_max(arr, arr.argmax(axis), axis) == arr.max(axis)"""
    new_shape = list(arr.shape)
    del new_shape[axis]

    grid = np.ogrid[tuple(map(slice, new_shape))]
    grid.insert(axis, argmax)

    return arr[tuple(grid)]

Desafortunadamente, un poco más incómodo de lo que debería ser una operación tan natural.

Para indexar una matriz n dim con una matriz n dim (n-1) dim , podríamos simplificarlo un poco para darnos la cuadrícula de índices para todos los ejes, así:

def all_idx(idx, axis):
    grid = np.ogrid[tuple(map(slice, idx.shape))]
    grid.insert(axis, idx)
    return tuple(grid)

Por lo tanto, úselo para indexar en matrices de entrada:

axis = 0
a_max_values = a[all_idx(idx, axis=axis)]
b_max_values = b[all_idx(idx, axis=axis)]