2

I have one 3D data = NxMxD numpy array, and another 2D idx = NxM integer array which values are in the range of [0, D-1]. I want to perform basic updates to each data = NxM entry at the depth given by the idx array at that position.

For example, for N = M = D = 2:

data = np.zeros((2,2,2))
idx = np.array([[0,0],[1, 1]], int)

And I want to perform a simple operation like:

data[..., idx] += 1

My expected output would be:

>>> data
array([[[ 1.,  0.],
        [ 1.,  0.]],

       [[ 0.,  1.],
        [ 0.,  1.]]])

idx indicates for each 2D coordinate which D should be updated. The above operation doesn't work.

I've found this approach in SO which solves the indexing problem by using:

data[np.arange(N)[:, None], np.arange(M)[None, :], idx] += 1

It works fine, but looks pretty horrible needing to manually index the whole matrices for what it seems a pretty simple operation (using one matrix as a index mask for the last channel).

Is there any better solution?

Imanol Luengo
  • 15,366
  • 2
  • 49
  • 67

1 Answers1

2

With numpy.ix_ it does not look so horrible but the underlying idea using fancy indexing is still the same

x = np.arange(N)
y = np.arange(M)

xx,yy = np.ix_(x,y)

data[xx,yy,idx] += 1

Note

The problem is that you want to change the values of data. If you just wanted to have the values according to idx you could do

out = np.choose(idx,data.transform(2,0,1))

However, this gives you a copy of the values of data and not a view which means that 

out += 1

has no effect on your values in data.

plonser
  • 3,323
  • 2
  • 18
  • 22