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로 후행 축을 바꿀 및 축 = -1로 FUNC을 적용 할 수

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