Partially discrete colormap matplotlib

Question

While there are a number of examples of discrete colormap (a,b,c), I would like to do something a little different. I want to have a 3D surface plot that has a sharp contrast between a small value and zero, so the colors 'jump' or the colormap is partially discrete. My reason for this is that I want to more clearly distinguish between small values and what is consider to be 'zero' within a plot.

I am generating a 3D surface plot and want to use a colormap (like 'terrain') to indicate height on the Z-axis. However, I want there to be a 'gap' in the colormap to highlight values that are sufficiently far from z=0. Specifically, let's say z<1e-6 is the bottom threshold of the colormap (e.g., dark blue for terrain), but any value above that threshold to be in the middle of the colormap (e.g. green for terrain).

Below is a simple example and the corresponding output

import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt

y = np.linspace(-3, 3, 100)
x = np.linspace(-3, 3, 100)
z = np.zeros(shape=(x.shape[0], y.shape[0]))

for i in range(x.shape[0]):
    # creating some generic Z-axis data
    z[:, i] = norm.pdf(y, loc=0, scale=0.2+(i/100))
    z[:, i] = z[:, i] / np.sum(z[:, i])  # normalizing

z = np.where(z < 1e-6, 0, z)  # setting 'small enough' threshold

x_mat, y_mat = np.meshgrid(x, y)

f1 = plt.axes(projection='3d')
f1.plot_surface(x_mat, y_mat, z, cmap='terrain', edgecolor='none', rstride=1)
plt.show()

Here is what the output from above:

What I want the output to look like would be all the 'light blue' regions would instead be green. Once below the defined threshold (1e-6 here), the color would jump to dark blue (so no regions would be light blue).

You may also look at the docs for creating your own colormaps: https://matplotlib.org/stable/tutorials/colors/colormap-manipulation.html — Jody Klymak, Jul 24 '21 at 14:33

score 0 · Accepted Answer · answered Jul 24 '21 at 14:11

Alright, I figured out a solution to my own problem. I adapted the solution from HERE to address my issue. Below is the code to accomplish this.

Setup:

import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap

y = np.linspace(-3, 3, 100)
x = np.linspace(-3, 3, 100)
z = np.zeros(shape=(x.shape[0], y.shape[0]))

x_mat, y_mat = np.meshgrid(x, y)

# Threshold to apply
z_threshold = 1e-6


for i in range(x.shape[0]):
    z[:, i] = norm.pdf(y, loc=0, scale=0.2+(i/100))
    z[:, i] = z[:, i] / np.sum(z[:, i])  # normalizing

Next I define two different colormaps. The first color map applies to all values above the threshold. If values are below the threshold, it sets that square as transparent.

cmap = get_cmap('terrain')
# 1.8 and 0.2 are used to restrict the upper and lower parts of the colormap
colors = cmap((z - z_threshold) / ((np.max(z)*1.8) - (np.min(z))) + 0.2)
# if below threshold, set as transparent (alpha=0)
colors[z < z_threshold, -1] = 0

The second colormap defines the color for all places below the threshold. This step isn't fully necessary, but it does prevent the plane from being drawn below the rest of the plot.

colors2 = cmap(z)
colors2[z >= z_threshold, -1] = 0

Now the colormaps can be used in two 3D plot calls

# init 3D plot
f1 = plt.axes(projection='3d')
# Plot values above the threshold
f1.plot_surface(x_mat, y_mat, z, facecolors=colors, edgecolor='none', rstride=1)
# Plot values below the threshold
z_below = np.zeros(shape=(x.shape[0], y.shape[0]))
f1.plot_surface(x_mat, y_mat, z_below,
                facecolors=colors2, edgecolor='none', rstride=1, vmin=0)
# Setting the zlimits
f1.set_zlim([0, np.max(z)])
plt.show()

The above results in the following plot

Partially discrete colormap matplotlib

1 Answers1