NumPy中的Fancier Fancy Indexing？

Question

I am implementing color interpolation using a look-up-table (LUT) with NumPy. At one point I am using the 4 most significant bits of RGB values to choose corresponding CMYK values from a 17x17x17x4 LUT. Right now it looks something like this:

import numpy as np
rgb = np.random.randint(16, size=(3, 1000, 1000))
lut = np.random.randint(256, size=(17, 17, 17, 4))
cmyk = lut[rgb[0], rgb[1], rgb[2]]

Here comes the first question... Is there no better way? It sort of seems natural that you could tell NumPy that the indices for lut are stored along axis 0 of rgb , without having to actually write it out. So is there anything like cmyk = lut.fancier_take(rgb, axis=0) in NumPy?

Furthermore, I am left with an array of shape (1000, 1000, 4) , so to be consistent with the input, I need to rotate it all around using a couple of swapaxes :

cmyk = cmyk.swapaxes(2, 1).swapaxes(1, 0).copy()

And I also need to add the copy statement, because if not the resulting array is not contiguous in memory, and that brings trouble later on.

Right now I am leaning towards rotating the LUT before the fancy indexing and then do something along the lines of:

swapped_lut = lut.swapaxes(2, 1).swapaxes(1, 0)
cmyk = swapped_lut[np.arange(4), rgb[0], rgb[1], rgb[2]]

But again, it just does not seem right... There has to be a more elegant way to do this, right? Something like cmyk = lut.even_fancier_take(rgb, in_axis=0, out_axis=0) ...

Answer 1

I'd suggest using tuple to force indexing rowwise, and np.rollaxis or transpose instead of swapaxes :

lut[tuple(rgb)].transpose(2, 0, 1).copy()

or

np.rollaxis(lut[tuple(rgb)], 2).copy()

To roll the axis first, use:

np.rollaxis(lut, -1)[(Ellipsis,) + tuple(rgb)]

Answer 2

You'll need to do the following if you swap lut , np.arange(4) will not work:

swapped_lut = np.rollaxis(lut, -1)
cmyk = swapped_lut[:, rgb[0], rgb[1], rgb[2]].copy()

Or you can replace

cmyk = lut[rgb[0], rgb[1], rgb[2]]
cmyk = cmyk.swapaxes(2, 1).swapaxes(1, 0).copy()

with:

cmyk = lut[tuple(rgb)]
cmyk = np.rollaxis(cmyk, -1).copy()

But to try and do it all in one step, ... Maybe:

rng = np.arange(4).reshape(4, 1, 1)
cmyk = lut[rgb[0], rgb[1], rgb[2], rng]

That's not very readable at all is it?

Take a look at the answer to this question, Numpy multi-dimensional array indexing swaps axis order . It does a good job of explaining how numpy broadcasts multiple arrays to get the output size. Here you want to create indices into lut that broadcast to (4, 1000, 1000). Hope that makes some sense.