Как перебрать "некоторые" измерения массива numpy?

Question

У меня есть трехмерный массив фигур (2,2,2). Я хотел бы думать об этом как о двумерной матрице с одномерными массивами в качестве записей.

Я хотел бы применить функцию к каждой записи (то есть к каждому 1-му массиву) в моей матрице. Я знаю, что могу векторизовать свою функцию, чтобы применить ее к каждому числу в моем массиве. Я также знаю, что могу применить функцию вдоль одной оси. Однако мне не удалось применить функции по двум осям.

Вот мое последнее испытание:

import numpy as np

def sqrtSum(a, b):
    return np.sqrt(a+b)

def sqrtSumWrapper(row):
    return np.array([sqrtSum(x[0], x[1]) for x in row])

z = np.array([[[1,2],[3,4]],[[5,6],[7,8]]])
np.apply_along_axis(sqrtSumWrapper, 1, z)

В приведенном выше примере моим желаемым результатом будет массив формы (2,2) с записями ((sqrt (3), sqrt (7)), (sqrt (11), sqrt (15)). Когда я запускаю приведенный выше код, я получаю недопустимый индекс для скалярной переменной. error.

Я думаю, что мне не хватает важного аспекта работы apply_along_axis и буду благодарен за ваш вклад в то, как исправить код.

EDIT: Пока ответы направлены на изменение функции ввода sqrtSum . Эта функция является лишь примером. Меня интересует общий ответ для произвольной входной функции, которая принимает n входных параметров и возвращает скаляр.

Решение : Решение удивительно простое (позор, что я его не видел)

import numpy as np
def sqrtSum(a, b):
    return np.sqrt(a+b)

def sqrtSumWrapper(x):
    return sqrtSum(x[0], x[1])

z = np.array([[[1,2],[3,4]],[[5,6],[7,8]]])
np.apply_along_axis(sqrtSumWrapper, 2, z)

Спасибо всем ответчикам.

Answer 1

Вот как правильно использовать numpy.apply_along_axis. Применяемая функция должна быть функцией одномерного массива:

def sqrtSum(arr):
    return np.sqrt(np.sum(arr))

>>> z = np.array([[[1,2],[3,4]],[[5,6],[7,8]]])
>>> np.apply_along_axis(sqrtSum, 2, z)
array([[1.73205081, 2.64575131],
       [3.31662479, 3.87298335]])

Для сравнения:

>>> np.array([[np.sqrt(3), np.sqrt(7)],[np.sqrt(11), np.sqrt(15)]])
array([[1.73205081, 2.64575131],
       [3.31662479, 3.87298335]])

Answer 2

вы можете использовать функции numpy's sum и sqrt.

z = np.array([[[1,2],[3,4]],[[5,6],[7,8]]])
s = np.sum(z,axis = 2)
print s.tolist()
# [[3, 7], [11, 15]]
sqrt = np.sqrt(s)