NumPy：將索引選擇功能應用於多個軸

Question

我正在編寫代碼以優化取決於可變數量參數的數量。爲了優化，我希望一次性在多個軸上應用索引選擇功能，例如numpy.argmax和numpy.argmin。以下是我現在使用的代碼。是否有更多的內置或有效的方法來執行跨任意數量的軸可能或不可能順序執行此任務？

def nd_arg_axes(func, array, start): 
    """Applies an index selecting function over trailing axes from start.""" 

    n_trail = len(array.shape[start:]) # Number of trailing axes to apply to. 

    indices = np.zeros((n_trail,)+array.shape[:start], dtype=np.intp) 
    for i in np.ndindex(array.shape[:start]): 
     indices[(Ellipsis,)+i] = np.unravel_index(func(array[i]), 
               array.shape[start:]) 
    return tuple(indices) 

# Test showing nd_arg_axes does indeed return the correct indices. 
array = np.arange(27).reshape(3,3,3) 
max_js = nd_arg_axes(np.argmax, array, 1) 

(array[tuple(np.indices(array))+max_js] == 
np.squeeze(np.apply_over_axes(np.amax, array, axes=[1,2]))) 

Answer 1

如果要選擇在軸的聯動軸，可以重塑後的軸爲-1，並應用到FUNC軸= -1：

def f(func, array, start): 
    shape = array.shape 
    tmp = array.reshape(shape[:start] + (-1,)) 
    indices = func(tmp, axis=-1) 
    return np.unravel_index(indices, shape[start:])