How do I apply function to third-dimension array effectively with numpy?
<p>I want to apply arbitrary function to 3d-ndarray as element, which use (3rd-dimensional) array for its arguments and return scalar.As a result, we should get 2d-Matrix.</p>

<p>e.g) pseudo code </p>

<pre><code>A = [[[1,2,3],[4,5,6]],
     [[7,8,9],[10,11,12]]]
A.apply_3d_array(sum) ## or apply_3d_array(A,sum) is Okey.
&gt;&gt; [[6,15],[24,33]]
</code></pre>

<p>I understand it's possible with loop using ndarray.shape function,but direct index access is inefficient as official document says. 
Is there more effective way than using loop? </p>

<pre><code>def chromaticity(pixel):
    geo_mean = math.pow(sum(pixel),1/3)
    return map(lambda x: math.log(x/geo_mean),pixel ) 
</code></pre>

<p>Given the function implementation, we could vectorize it using <a href="http://docs.scipy.org/doc/numpy/reference/ufuncs.html" rel="nofollow"><code>NumPy ufuncs</code></a> that would operate on the entire input array <code>A</code> in one go and thus avoid the <code>math</code> library functions that doesn't support vectorization on arrays.  In this process, we would also bring in the very efficient vectorizing tool : <a href="http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html" rel="nofollow"><code>NumPy broadcasting</code></a>. So, we would have an implementation like so -</p>

<pre><code>np.log(A/np.power(np.sum(A,2,keepdims=True),1/3))
</code></pre>

<p><strong>Sample run and verification</strong></p>

<p>The function implementation without the <code>lamdba</code> construct and introducing NumPy functions instead of <code>math</code> library functions, would look something like this -</p>

<pre><code>def chromaticity(pixel): 
    geo_mean = np.power(np.sum(pixel),1/3) 
    return np.log(pixel/geo_mean)
</code></pre>

<p>Sample run with the iterative implementation -</p>

<pre><code>In [67]: chromaticity(A[0,0,:])
Out[67]: array([-0.59725316,  0.09589402,  0.50135913])

In [68]: chromaticity(A[0,1,:])
Out[68]: array([ 0.48361096,  0.70675451,  0.88907607])

In [69]: chromaticity(A[1,0,:])
Out[69]: array([ 0.88655887,  1.02009026,  1.1378733 ])

In [70]: chromaticity(A[1,1,:])
Out[70]: array([ 1.13708257,  1.23239275,  1.31940413])    
</code></pre>

<p>Sample run with the proposed vectorized implementation -</p>

<pre><code>In [72]: np.log(A/np.power(np.sum(A,2,keepdims=True),1/3))
Out[72]: 
array([[[-0.59725316,  0.09589402,  0.50135913],
        [ 0.48361096,  0.70675451,  0.88907607]],

       [[ 0.88655887,  1.02009026,  1.1378733 ],
        [ 1.13708257,  1.23239275,  1.31940413]]])
</code></pre>

<p><strong>Runtime test</strong></p>

<pre><code>In [131]: A = np.random.randint(0,255,(512,512,3)) # 512x512 colored image

In [132]: def org_app(A):
     ...:     out = np.zeros(A.shape)     
     ...:     for i in range(A.shape[0]):
     ...:         for j in range(A.shape[1]):
     ...:             out[i,j] = chromaticity(A[i,j])
     ...:     return out
     ...: 

In [133]: %timeit org_app(A)
1 loop, best of 3: 5.99 s per loop

In [134]: %timeit np.apply_along_axis(chromaticity, 2, A) #@hpaulj's soln
1 loop, best of 3: 9.68 s per loop

In [135]: %timeit np.log(A/np.power(np.sum(A,2,keepdims=True),1/3))
10 loops, best of 3: 90.8 ms per loop
</code></pre>

<p>That's why always try to push in <code>NumPy funcs</code> when vectorizing things with arrays and work on as many elements in one-go as possible!</p>

python arrays numpy multidimensional-array vectorization