Color numpy 2-dim array based on condition

Question

I have 2d array with values between -0.5 to 1.5 that looks like this:

I would like to color this array with two colors: if values are lower than 0, I would like it to be green (but not one "shade" of green,more like another cmap of green), and if is greater than 0 , to have red color (again. not one specific color, but more like cmap with ranges). So the result should look something like this (conceptually):

I haven't found yet way to apply these "two cmaps" , so my end goal is to color this one array into two colors, based on the conditions.

Edit: example:

np.array([0.1,0.15,0.2,0.5,0.7,
         -0.1,-0.12,0.23,0.4,0.8,
         -0.01,-.15,-0.1,0.2,0.5]).reshape(5,3)

Answer 1

Using different colours for different ranges of your data is the job of a colourmap, so if I were you I'd try to avoid having to plot the data twice (i.e. overlapping plots with two different colourmaps, which seems like what you were maybe thinking). Then you can either choose an existing colourmap, or make one... or I can think of one other way to do it.

Note that in the examples here I'm using slightly different data to the array you gave.

Choosing an existing 'diverging' colourmap is straightforward, just be sure to center it where you want to make your cutoff, so to speak. In your case, that's zero:

import numpy as np
import matplotlib.pyplot as plt

arr = np.array([ 0.1,   0.15,  0.2,  0.5, 0.7,
                -0.1,  -0.12,  0.23, 0.4, 0.8,
                -0.01, -0.45, -0.1,  0.2, 0.5]).reshape(3, 5)

plt.imshow(arr, cmap="PiYG_r", vmin=-1.5, vmax=1.5)  # Center it.
plt.colorbar(shrink=0.75)

This yields:

If you want a custom colourmap, you can interpolate one using LinearSegmentedColormap. This is probably what I would do in your situation. Use it like so:

from matplotlib.colors import LinearSegmentedColormap

cdict = {'red':   [[0.0, 0.0, 0.0],  # Anchor, left, right.
                   [0.5, 0.0, 0.0],  # Black at centre point.
                   [1.0, 1.0, 1.0]],
         'green': [[0.0, 1.0, 1.0],
                   [0.5, 0.0, 0.0],
                   [1.0, 0.2, 0.2]],
         'blue':  [[0.0, 0.0, 0.0],
                   [0.5, 0.0, 0.0],
                   [1.0, 0.4, 0.4]]}

newcmp = LinearSegmentedColormap('newCmap', segmentdata=cdict, N=256)
plt.imshow(arr, cmap=newcmp, vmin=-1.5, vmax=1.5)
plt.colorbar(shrink=0.75)

Which results in:

Alternatively, you could compute a blended RGB image by building up an image array like this:

red = arr.copy()
red[arr < 0] = np.nan

grn = arr.copy()
grn[arr >= 0]  = np.nan

blu = np.zeros_like(red)*np.nan

rgb = np.dstack([red, np.abs(grn), blu])
plt.imshow(rgb)

This produces:

The main drawback here is that there is no colorbar, as such (it's an RGB cube).