Какой алгоритм стоит за опцией оси в numpy?

Question

Я не могу понять, как ось в numpy работает для любого общего массива nD .

Например, np.mean(array, axis=1) даст ожидаемый массив независимо от формы array.

Если я хочу создать код без использования numpy, я могу сделать функцию усреднения для массива kD, например, mean_1d, mean_2d, ... с опцией оси. Однако я не могу представить, как создать единственную функцию, которая работает для любого входного массива с любым значением k. Пока не нашел объяснения, как работает ось numpy. Может ли кто-нибудь помочь мне понять, что происходит за сценой? массивы и получить (n-1) -D массив) с помощью numba.

Answer 1

Если размер ваших массивов позволяет, вы можете увидеть это рекурсивно. Я только что протестировал приведенный ниже код, и, похоже, он отлично работает для функции суммы вдоль оси 3. Это просто, но должно передать идею:

import numpy as np

a = np.ones((3,4,5,6,7))

# first get the shape of the result
def shape_res(sh1, axis) :
    sh2 = [sh1[0]]

    if len(sh1)>0 :
        for i in range(len(sh1)) :
            if i > 0 and i != axis :
                sh2.append(sh1[i])
            elif i == axis :
                # shrink the summing axis
                sh2.append(1)
    return sh2


def sum_axis(a, axis=0) : 
    cur_sum = 0
    sh1 = np.shape(a)
    sh2 = shape_res(sh1, axis)

    res = np.zeros(sh2)
    print(sh1, sh2)

    def rec_fun (a, res, cur_axis) :
        for i in range(len(a)) :
            if axis > cur_axis :
                # dig in the array to reach the right axis
                next_axis = cur_axis + 1
                rec_fun(a[i], res[i], next_axis)
            elif axis == cur_axis :
                # sum along the given axis
                res[0] += a[i]

    rec_fun(a, res, cur_axis)
    return res

result = sum_axis(a, axis=3)
print(result)