CODEWITHSUNDEEP
pythonnumpycolorbarimshow
Reut
Reut

Reputation: 1592

Color numpy 2-dim array based on condition

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

enter image description here

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): enter image description here

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)

Upvotes: 2

Views: 632

Answers (1)

Matt Hall
Matt Hall

Reputation: 8132

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:

Result of using PiYG_r

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:

Result of using custom colourmap

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:

Result of using RGB blend

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

Upvotes: 1

Related Questions