title: Apply a function to subsample of 3-dimensional numpy array - variation on np.apply_along_axis
body: <p>I have a numpy array of 3D vectors of shape <code>(n, 3)</code>, n being quite large, which can be illustrated as follows:</p>
  <pre><code>vects = [[x_1  y_1  z_1]
           ...
           [x_i  y_i  z_i]
           ...
           [x_n  y_n  z_n]]
  </code></pre>
  <p>I then use an existing function computing the divergence of my vectors array according to spatial coordinates <code>DataSet</code> at witch the vectors are specified:</p>
  <p><code>vects_div = algs.divergence(vects, DataSet)</code></p>
  <p>What is computed actually doesn't really matters, the important point is that my function needs as arguments:</p>
  <ul>
  <li><strong>a <code>(n, 3)</code> shaped numpy array</strong></li>
  <li>an object describing the spatial coordinates</li>
  </ul>
  <p>and is outputting a <code>(n,)</code> array of scalars.</p>
  <p><em>My issue is that I want to apply this function to subsamples of an array of tensors.</em></p>
  <p>Let me explain.</p>
  <p>This time, I have a numpy array of 3D tensors of shape <code>(n, 3, 3)</code>, which can be illustrated as follows:</p>
  <pre><code>tensors = [[[ xx_1  xy_1  xz_1]
              [ yx_1  yy_1  yz_1]
              [ zx_1  zy_1  zz_1]]
            ...
             [[ xx_i  xy_i  xz_i]
              [ yx_i  yy_i  yz_i]
              [ zx_i  zy_i  zz_i]]
            ...
             [[ xx_n  xy_n  xz_n]
              [ yx_n  yy_n  yz_n]
              [ zx_n  zy_n  zz_n]]]
  </code></pre>
  <p>I then use the previous function on three vectors arrays subsamples of this tensors array:</p>
  <pre><code>tensors_x = tensors[:,0,:]
  tensors_x = [[xx_1  xy_1  xz_1]
               ...
               [xx_i  xy_i  xz_i]
               ...
               [xx_n  xy_n  xz_n]]
  
  tensors_y = tensors[:,1,:]
  tensors_y = [[yx_1  yy_1  yz_1]
               ...
               [yx_i  yy_i  yz_i]
               ...
               [yx_n  yy_n  yz_n]]
  
  tensors_z = tensors[:,2,:]
  tensors_z = [[zx_1  zy_1  zz_1]
               ...
               [zx_i  zy_i  zz_i]
               ...
               [zx_n  zy_n  zz_n]]
  </code></pre>
  <pre><code>tensors_x_div = algs.divergence(tensors_x, DataSet)
  
  tensors_y_div = algs.divergence(tensors_y, DataSet)
  
  tensors_z_div = algs.divergence(tensors_z, DataSet)
  </code></pre>
  <p>Which results in three separate arrays of scalars of shape <code>(n,)</code>.</p>
  <p>Finally I rebuild the vectors array of results as follows:</p>
  <pre><code>tensors_div = np.column_stack((tensors_x_div, tensors_y_div, tensors_z_div))
  </code></pre>
  <p>All of this can be summarized into a single and dirty line:</p>
  <pre><code>tensors_div = np.column_stack((algs.divergence(tensors[:,0,:], DataSet), algs.divergence(tensors[:,1,:], DataSet), algs.divergence(tensors[:,2,:], DataSet)))
  
  </code></pre>
  <p>This is all working fine, but I'am looking for a prettier and if possible more efficient way to do it.</p>
  <p>I was thinking at the <code>np.apply_along_axis</code> function with something like</p>
  <pre><code>np.apply_along_axis(algs.divergence, 1, tensors, DataSet)
  </code></pre>
  <p>But the slices sent to my function are of shape <code>(3,)</code> instead of <code>(n,3)</code>.</p>
  <p>Is there any numpy function which would allow to do what I want ?</p>
  <p>Is there a better way than my dirty line ?</p>