title: Using numpy.take for faster fancy indexing
body: <p><strong>EDIT</strong> I have kept the more complicated problem I am facing below, but my problems with <code>np.take</code> can be summarized better as follows. Say you have an array <code>img</code> of shape <code>(planes, rows)</code>, and another array <code>lut</code> of shape <code>(planes, 256)</code>, and you want to use them to create a new array <code>out</code> of shape <code>(planes, rows)</code>, where <code>out[p,j] = lut[p, img[p, j]]</code>. This can be achieved with fancy indexing as follows:</p>
  
  <pre><code>In [4]: %timeit lut[np.arange(planes).reshape(-1, 1), img]
  1000 loops, best of 3: 471 us per loop
  </code></pre>
  
  <p>But if, instead of fancy indexing, you use take and a python loop over the <code>planes</code> things can be sped up tremendously:</p>
  
  <pre><code>In [6]: %timeit for _ in (lut[j].take(img[j]) for j in xrange(planes)) : pass
  10000 loops, best of 3: 59 us per loop
  </code></pre>
  
  <p>Can <code>lut</code> and <code>img</code> be in someway rearranged, so as to have the whole operation happen without python loops, but using <code>numpy.take</code> (or an alternative method) instead of conventional fancy indexing to keep the speed advantage?</p>
  
  <hr>
  
  <p><strong>ORIGINAL QUESTION</strong>
  I have a set of look-up tables (LUTs) that I want to use on an image. The array holding the LUTs is of shape <code>(planes, 256, n)</code>, and the image has shape <code>(planes, rows, cols)</code>. Both are of <code>dtype = 'uint8'</code>, matching the <code>256</code> axis of the LUT. The idea is to run the <code>p</code>-th plane of the image through each of the <code>n</code> LUTs from the <code>p</code>-th plane of the LUT.</p>
  
  <p>If my <code>lut</code> and <code>img</code> are the following:</p>
  
  <pre><code>planes, rows, cols, n = 3, 4000, 4000, 4
  lut = np.random.randint(-2**31, 2**31 - 1,
                          size=(planes * 256 * n // 4,)).view('uint8')
  lut = lut.reshape(planes, 256, n)
  img = np.random.randint(-2**31, 2**31 - 1,
                      size=(planes * rows * cols // 4,)).view('uint8')
  img = img.reshape(planes, rows, cols)
  </code></pre>
  
  <p>I can achieve what I am after using fancy indexing like this</p>
  
  <pre><code>out = lut[np.arange(planes).reshape(-1, 1, 1), img]
  </code></pre>
  
  <p>which gives me an array of shape <code>(planes, rows, cols, n)</code> , where <code>out[i, :, :, j]</code> holds the <code>i</code>-th plane of <code>img</code> run through the <code>j</code>-th LUT of the <code>i</code>-th plane of the LUT...</p>
  
  <p>All is good, except for this:</p>
  
  <pre><code>In [2]: %timeit lut[np.arange(planes).reshape(-1, 1, 1), img]
  1 loops, best of 3: 5.65 s per loop
  </code></pre>
  
  <p>which is completely unacceptable, especially since I have all of the following not so nice looking alternatives using <code>np.take</code> than run much faster:</p>
  
  <ol>
  <li><p>A single LUT on a single plane runs about x70 faster:</p>
  
  <pre><code>In [2]: %timeit np.take(lut[0, :, 0], img[0])
  10 loops, best of 3: 78.5 ms per loop
  </code></pre></li>
  <li><p>A python loop running through all the desired combinations finishes almost x6 faster:</p>
  
  <pre><code>In [2]: %timeit for _ in (np.take(lut[j, :, k], img[j]) for j in xrange(planes) for k in xrange(n)) : pass
  1 loops, best of 3: 947 ms per loop
  </code></pre></li>
  <li><p>Even running all combinations of planes in the LUT and image and then discarding the <code>planes**2 - planes</code> unwanted ones is faster than fancy indexing:</p>
  
  <pre><code>In [2]: %timeit np.take(lut, img, axis=1)[np.arange(planes), np.arange(planes)]
  1 loops, best of 3: 3.79 s per loop
  </code></pre></li>
  <li><p>And the fastest combination I have been able to come up with has a python loop iterating over the planes and finishes x13 faster: </p>
  
  <pre><code>In [2]: %timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass
  1 loops, best of 3: 434 ms per loop
  </code></pre></li>
  </ol>
  
  <p>The question of course is if there is no way of doing this with <code>np.take</code> without any python loop? Ideally whatever reshaping or resizing is needed should happen on the LUT, not the image, but I am open to whatever you people can come up with...</p>