Use lambda with multi input with numpy.apply_along_axis

Here is my code:

# 'a' is a 3D array which is the RGB data
a = np.array([[[1, 2, 3], [4, 5, 6]], [[4, 5, 6], [1, 2, 3]]])
# I want to do some calculate separately on R, G and B
np.apply_along_axis(lambda r, g, b: r * 0.5 + g * 0.25 + b * 0.25, axis=-1, arr=a)

# my target output is: [[1.75, 4.75], [4.75, 1.75]]

But the above way will give me error, missing 2 required positional arguments.

I have tried to do like this:

np.apply_along_axis(lambda x: x[0] * 0.5 + x[1] * 0.25 + x[2] * 0.25, axis=-1, arr=a)

It works but every time when I need to do computation on the array element, I need to type the index, it is quite redundant. Is there any way that I can pass the array axis as multi iuputs to the lambda when using np.apply_along_axis?

Answer

apply_along_axis is slow and unnecessary:

In [279]: a = np.array([[[1, 2, 3], [4, 5, 6]], [[4, 5, 6], [1, 2, 3]]])
In [280]: r,g,b = a[...,0],a[...,1],a[...,2]
     ...: r * 0.5 + g * 0.25 + b * 0.25
Out[280]: 
array([[1.75, 4.75],
       [4.75, 1.75]])

or

In [281]: a.dot([0.5, 0.25, 0.25])
Out[281]: 
array([[1.75, 4.75],
       [4.75, 1.75]])

or

In [282]: np.sum(a*[0.5, 0.25, 0.25], axis=2)
Out[282]: 
array([[1.75, 4.75],
       [4.75, 1.75]])