How to efficiently match two array values in Python?

Question

I am trying to find the best way to match two array axis in Python, most likely using Numpy. To be more precise : I have C which is a H x W x L matrix and D which is a H x W matrix. H and W are the height and width of an image, and L is, for example, the set of luminance that can be displayed in the image ( 0 to 100 for example). D contains the real luminance of the pixel.

I want to add all the values linked to the luminance in the third axis of C of the pixels that match the luminance in the D matrix, eg the value in the D matrix would tell me which index to use to get the value in the C matrix.

Of course, I can do this easily with three loops :

sum = 0

for row in range(H) :
  for column in range(W) :
    for luminance in range(L) :
      if luminance == D[row, column] :
        sum += C[row, column, luminance]
        break

But this is not efficient. I am trying to find a way to do this using numpy, something like this :

import numpy as np

sum = np.sum(C[:, :, np.where(C[:,:] == [D[:,:]])[0][0]))

I do not know if this is supposed to work. I think there might be a way using the meshgrid function to loop through the pixels easily, and correctly have the match by calling

import numpy as np

sum = np.sum(np.meshgrid(range(H), range(W), f(D)))

where f(D) has to be something to do with the D matrix.

Well, this is what I tried so far, thank you for your help as always !

Answer 1

Why don't you just sum all values of D ? Is there a chance that the luminance in C won't appear in D , is that why?

Anyway, another way of accomplishing what you want is by repeating values of C to match the shape of D , and then summing only the values that match, provided that each pixel of C doesn't have duplicate values.

D_repeated = np.repeat(D[:, :, np.newaxis], L, axis=2)
sum = np.sum(C[C==D_repeated])

Answer 2

I am bumping this thread to include an answer using numpy.meshgrid .

One would be able to do that with :

import numpy as np

lines, columns = np.meshgrid(range(H), range(W), indexing='ij')

sum = np.sum(C[lines, columns, D])

With that, the value in D tells which index to use in C to get the right value.