CODEWITHSUNDEEP
pythonnumpy
Haddock
Haddock

Reputation: 3

Apply function to vectors in 3D numpy array

I have a question about how to apply a function to vectors in a 3D numpy array. My problem is the following: let's say I have an array like this one:

a = np.arange(24)
a = a.reshape([4,3,2])

I want to apply a function to all following vectors to modify them:

[0 6], [1 7], [2 8], [4 10], [3 9] ...

What is the best method to use? As my array is quite big, looping in two of the three dimension is quite long...

Thanks in advance!

Upvotes: 0

Views: 955

Answers (1)

RomanHotsiy
RomanHotsiy

Reputation: 5138

You can use function np.apply_along_axis. From the doc:

Apply a function to 1-D slices along the given axis.

For example:

>>> import numpy as np
>>> a = np.arange(24)
>>> a = a.reshape([4,3,2])
>>> 
>>> def my_func(a):
...   print "vector: " + str(a)
...   return sum(a) / len(a)
... 
>>> np.apply_along_axis(my_func, 0, a)
vector: [ 0  6 12 18]
vector: [ 1  7 13 19]
vector: [ 2  8 14 20]
vector: [ 3  9 15 21]
vector: [ 4 10 16 22]
vector: [ 5 11 17 23]
array([[ 9, 10],
       [11, 12],
       [13, 14]])

In example above I've used 0th axis. If you need n axes you can execute this function n times.

Upvotes: 1

Related Questions