How to iterate over "some" dimensions of a numpy array?

Question

I've got a three-dimensional numpy array of shape (2,2,2). I would like to think of it as a two-dimensional matrix with 1-dimensional arrays as entries.

What I would like to do is applying a function to each entry (i.e. each 1-d array) in my matrix. I know that I could vectorize my function to apply it to each number in my array. I also know that I can apply a function along one axis. However, I haven't managed to apply the functions along two axes.

Here's my latest trial:

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)

In the example above, my desired outcome would be an array of shape (2,2) with the entries ((sqrt(3), sqrt(7)),(sqrt(11), sqrt(15)). When I run the code above I get a invalid index to scalar variable. error.

I think I'm missing an important aspect of how apply_along_axis works and would be thankful for your input how to correct the code.

EDIT: The answers so far are focussing on changing the input function sqrtSum. This function is just an example. I'm interested in a general answer for an arbitrary input function that takes n input parameters and returns a scalar.

Solution: The solution is amazingly simple (shame on me that I haven't seen it)

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)

Thanks to all repliers.

Answer 1

Here's how to use numpy.apply_along_axis correctly. The function to apply must be a function of a 1D array:

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

 

For comparison:

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