如何使用一维的子索引生成多维2D numpy索引

Question

I want to use numpy.ix_ to generate an multi-dimensional index for a 2D space of values. However, I need to use a subindex to look up the indices for one dimension. For example,

    assert subindex.shape == (ny, nx)

    data = np.random.random(size=(ny,nx))

    # Generator returning the index tuples 
    def get_idx(ny,nx,subindex):
      for y in range(ny):
        for x in range(nx):
           yi = y             # This is easy
           xi = subindex[y,x] # Get the second index value from the subindex

           yield (yi,xi)

    # Generator returning the data values
    def get_data_vals(ny,nx,data,subindex):
      for y in range(ny):
        for x in range(nx):
           yi = y             # This is easy
           xi = subindex[y,x] # Get the second index value from the subindex

           yield data[y,subindex[y,x]]

So instead of the for loops above, I'd like to use a multi-dimensional index to index data Using numpy.ix_ , I guess I would have something like:

    idx = numpy.ix_([np.arange(ny), ?])
    data[idx]

but I don't know what the second dimension argument should be. I'm guessing it should be something involving numpy.choose ?

Answer 1

It sounds like you need to do two things:

• Find all indices into the data array and
• Translate the column indices according to some other array, subindex.

The code below therefore generates indices for all array positions (using np.indices ), and reshapes it to (..., 2) -- a 2-D list of coordinates representing each position in the array. For each coordinate, (i, j) , we then translate the column coordinate j using the subindex array provided, and then use that translated index as the new column index.

With numpy, it is not necessary to do that in a for-loop--we can simply pass in all the indices at once:

i, j = np.indices(data.shape).reshape((-1, 2)).T
data[i, subindex[i, j]]

Answer 2

What you actually seem to want is:

y_idx = np.arange(ny)[:,np.newaxis]
data[y_idx, subindex]

BTW, you could achieve the same thing with y_idx = np.arange(ny).reshape((-1, 1)) .

Let's look at a small example:

import numpy as np

ny, nx = 3, 5
data = np.random.rand(ny, nx)
subindex = np.random.randint(nx, size=(ny, nx))

Now

np.arange(ny)
# array([0, 1, 2])

are just the indices for the "y-axis", the first dimension of data . And

y_idx = np.arange(ny)[:,np.newaxis]
# array([[0],
#        [1],
#        [2]])

adds a new axis to this array (after the existing axis) and effectively transposes it. When you now use this array in an indexing expression together with the subindex array, the former gets broadcasted to the shape of the latter. So y_idx becomes effectively:

# array([[0, 0, 0, 0, 0],
#        [1, 1, 1, 1, 1],
#        [2, 2, 2, 2, 2]])

And now for each pair of y_idx and subindex you look up an element in the data array.

Here you can find out more about "fancy indexing"