Question: I have a 2d NumPy array with values from 0 to 1. I want to turn this array into a Pillow image. I can do the following, which gives me a nice greyscale image: arr = np.random.rand(100,100) img = Image.fromarray((255 * arr).astype(np.uint8))  Now, instead of making a greyscale image, I'd like to apply a custom gradient. To clarify, instead of drawing bands of colors in a linear gradient as in this example, I'd like to specify apply a gradient colormap to an existing 2d array and turn it into a 3d array. Example: If my gradient is [color1, color2, color3], then all 0s should be color1, all 1s should be color3, and 0.25 should be somewhere in between color1 and color2. I was already able to write a simple function that does this: gradient = [(0, 0, 0), (255, 80, 0), (0, 200, 255)] # black -> orange -> blue def get_color_at(x): assert 0 <= x <= 1 n = len(gradient) if x == 1: return gradient[-1] pos = x * (n - 1) idx1 = int(pos) idx2 = idx1 + 1 frac = pos - idx1 color1 = gradient[idx1] color2 = gradient[idx2] color_in_between = [round(color1[i] * (1 - frac) + color2[i] * frac) for i in range(3)] return tuple(color_in_between)  So get_color_at(0) returns (0,0,0) and get_color_at(0.75) equals (153, 128, 102), which is this tan/brownish color in between orange and blue. Now, how can I apply this to the original NumPy array? I shouldn't apply get_color_at directly to the NumPy array, since that would still give a 2d array, where each element is a 3-tuple. Instead, I think I want an array whose shape is (n, m, 3), so I can feed that to Pillow and create an RGB image. If possible, I'd prefer to use vectorized operations whenever possible - my input arrays are quite large. If there is builtin-functionality to use a custom gradient, I would also love to use that instead of my own get_color_at function, since my implementation is pretty naive. Thanks in advance.

Answer: Method 1: vectorization of your code Your code is almost already vectorized. Almost all operations of it can work indifferently on a float or on an array of floats Here is a vectorized version def get_color_atArr(arr): assert (arr>=0).all() and (arr<=1).all() n=len(gradient) gradient.append(gradient[-1]) gradient=np.array(gradient, dtype=np.uint8) pos = arr*(n-1) idx1 = pos.astype(np.uint8) idx2 = idx1+1 frac = (pos - idx1)[:,:,None] color1 = gradient[idx1] color2 = gradient[idx2] color_in_between = np.round(color1*(1-frac) + color2*frac).astype(np.uint8)  Basically, the changes are,  the assert (can't use a<b<c notation with numpy arrays). Note that this assert iterates all values of array to check for assertion. That is not for free. So I included it because you did. But you need to be aware that this is not a compile-time verification. It does run code to check all values, which is a non-negligible part of all execution time of the code. more an implementation choice than a vectorization step (a pure translation of your code would have translated that if x==1 into some np.where, or masks. But I am never comfortable with usage of == on floats any way. So I prefer my way. Which costs nothing. It is not another iteration on the image. It adds a sentinel (In Donald Kuth sense of "sentinel": a few bytes that avoid special cases) to the gradient color. So that, in the unlikely even that arr is really 1.0, the gradient happen between last color and last color). frac is broadcasted in 3D array, so that it can be used as a coefficient on 3d arrays color1 and color2 Plus of course, int or floor can't be used on numpy arrays  Method 2: not reinventing the wheel Matplotlib (and, I am certain, many other libraries) already have a whole colormap module to deal with this kind of transformations. Let's use it thresh=np.linspace(0,1,len(gradient)) cmap=LinearSegmentedColormap.from_list('mycmap', list(zip(thresh, np.array(gradient)/255.0)), N=256*len(gradient)) arr2 = cmap(arr)[:,:,:3]  This is building a custom colormap, using LinearSegmentedColormap, which takes, as 2nd argument, a list of pair (threshold, color). Such as [(0, (0,0,0)), (0.3, (1,0,0)), (0.8, (0,1,0)), (1, (0,0,1))] for a color map that goes from black to red when x goes from 0 tom 0.3, then from red to green when x goes from 0.3 to 0.8, then from green to blue. In this case, your gradient can be transformed to such a list, with just a zip with a linspace. It takes a N= argument, since it creates a discretization of all possible colors (with interpolation in between). Here I take an exaggerated option (my N is more than the maximum number of different colors than can exist, once uint8d) Also since it returns a RGBA array, and to remain strictly identical to what you did, I drop the A using [:,:,:3]. Of course, both method need the final translation into PIL, but you already know how to do that. For this one, it also needs mapping between 0 and 255, which I can do with your own code: Image.fromarray((255 * arr).astype(np.uint8)) Note that, while using matplotlib colormap, you may want to take a tour at what that module has to offer. For example some of the zillions of already existing colormaps may suit you. Or some other way to build colors map non-linearly.