CODEWITHSUNDEEP
pythonnumpy
Kenenbek Arzymatov
Kenenbek Arzymatov

Reputation: 9139

How to apply a function along an axis in numpy?

For example, I have a matrix with shape:

x = np.random.rand(3, 10, 2, 6)

As you can see, there are only two arrays along an axis=2.

I have a function that accepts these two arrays:

def f(arr1, arr2): # arr1 with shape (6, ) and arr2 with (6, )
    return np.sum(arr1, arr2) # for simplicity

How can I apply this function along the second axis to x array in a vectorized way? Such that resulting array will be of shape [3, 10, dim of output].

I came across apply_along_axis routine, but it requires that f accepts only 1D slice.

Upvotes: 1

Views: 717

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114518

You can't do it entirely arbitrarily, but your particular case reduces to

x.sum(axis=2)

If you want to add the arrays as in your code:

x[:, :, 0, :] + x[:, :, 1, :]

Upvotes: 1

Related Questions