indeksowi n Tablica wymiarową z (n-1) d tablicy

Question

co najbardziej elegancki sposób otwierającej n wymiarową tablicę z (n-1) Macierz wymiarów wzdłuż danego wymiaru jako zastępczej przykład

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

Jak mogę uzyskać dostęp teraz za pomocą idx a, aby uzyskać maksymalne wartości w a, tak jakbym używał a.max(axis=0)? lub jak odzyskać wartości określone przez idx w b?

Myślałem o użyciu np.meshgrid, ale myślę, że to przesada. Należy pamiętać, że wymiar axis może być dowolną użyteczną osią (0,1,2) i nie jest znany z góry. Czy istnieje elegancki sposób na zrobienie tego?

Answer 1

użycie Marka 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] 

---

W ogólnym przypadku:

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

Dość nieco bardziej niewygodne niż takiej naturalnej eksploatacji powinno być, niestety.

Dla indeksowania n dim tablicę z (n-1) dim tablicy, możemy go uprościć trochę dać nam siatkę indeksów dla wszystkich osi, tak jak -

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

Stąd też używać go do indeksu do tablic wejściowych -

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