如何迭代numpy數組的“某些”維度？

Question

我有一個三維的numpy形狀陣列（2,2,2）。 我想把它看作一個二維矩陣，一維數組作為條目。

我想要做的是在我的矩陣中為每個條目（即每個1-d數組）應用一個函數。 我知道我可以向量化我的函數以將其應用於我的數組中的每個數字。 我也知道我可以沿一個軸應用一個函數。 但是，我還沒有設法沿兩個軸應用這些功能。

這是我最新的試用版：

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））。當我運行上面的代碼我得到標量變量的無效索引。錯誤。

我想我錯過了apply_along_axis如何工作的一個重要方面，並感謝你輸入如何糾正代碼。

編輯：到目前為止的答案集中在改變輸入函數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的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)