[Numpy-discussion] 2D (or n-d) fancy indexing?

[Robert Kern]

On Wed, Oct 8, 2008 at 20:00, Zachary Pincus <zachary.pincus at yale.edu> wrote:
> Hello all,
>
> I'm doing something silly with images and am unable to figure out the
> right way to express this with "fancy indexing" -- or anything other
> than a brute for-loop for that matter.
>
> The basic gist is that I have an array representing n images, of shape
> (n, x, y). I also have a "map" of shape (x, y), which contains indices
> in the range [0, n-1]. I then want to construct the "composite" image,
> of shape (x, y), with pixels selected from the n source images as per
> the indices in the map, i.e.:
>
> composite[x, y] = images[map[x, y], x, y]
> for all (x, y).
>
> Now, I can't figure out if there's an easy way to express this in
> numpy. For that matter, I can't even figure out a simple way to do the
> 1D version of the same:
>
> composite[i] = images[map[i], i]
> where composite and map have shape (m,), and images has shape (n, m).
>
> Can anyone assist? Surely there's something simple that I'm just not
> seeing.

You need to give an array for each axis. Each of these arrays will be
broadcast against each other to form three arrays of the desired shape
of composite. This is discussed in the manual here:

  http://mentat.za.net/numpy/refguide/indexing.xhtml#indexing-multi-dimensional-arrays

Conceptually, you need arrays A, B, and C such that

  composite[x,y] == images[A[x,y], B[x,y], C[x,y]]
  for all x,y

You already have A, you just need to construct B and C correctly.

Here is an example:

In [26]: from numpy import *

In [27]: Nx = 480

In [28]: Ny = 640

In [29]: N = 100

In [30]: images = random.randint(0, 1000, size=(N, Nx, Ny))

In [31]: map = random.randint(0, N, size=(Nx, Ny))

In [32]: composite = images[map, arange(Nx)[:,newaxis], arange(Ny)]

In [33]: for x in range(Nx):
   ....:     for y in range(Ny):
   ....:         assert composite[x,y] == images[map[x,y],x,y]
   ....:
   ....:

In [34]:


When arange(Nx)[:,newaxis] and arange(Ny) get broadcasted with map,
you get (480,640) arrays like you would get with mgrid[0:Nx,0:Ny].

-- 
Robert Kern

"I have come to believe that the whole world is an enigma, a harmless
enigma that is made terrible by our own mad attempt to interpret it as
though it had an underlying truth."
  -- Umberto Eco