I am trying to obtain a 2D pseudocolor plot with a nonlinear color map (cmap). Independently I want to have a colobar that uses a similar cmap but differently scaled/stretched to avoid overlapping of the colorbar yticks.
The first one I can obtain using some nonlinear norm as an argument of pcolormesh.
But how to get the second part in an efficient way?
Finally, I was able to obtain the desired effect (see the bottom right corner in the below figure) but I am pretty sure that this is not the best/easiest/desired/Pythonic way of doing it.
Is there an easier way of obtaining such an effect?
Figure:
Here is the code that reproduces the above figure:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cbook as cbook
from matplotlib import cm
import copy
def transform_cmap(cmap_basic,
new_cmap_name = None,
cmap_trans_fun = None,
how_many_levels = 256,
ticks = None,
ticks_trans_fun = None):
'''returns new cmap and new ticks locators for transformed ticks.
If ticks_trans_fun is None then ticks locators are linearly transformed
to the [0,1] interval, as a result, cmap is stretched, but when used with
colorbar than linearly spaced ticks are linearly spaced along the colorbar'''
# be sure that cmap is really a cmap
if not isinstance(cmap_basic,colors.Colormap):
try:
cmap_basic = cm.get_cmap(name=cmap_basic)
except:
print('basic_cmap is not a valid cmap or cmap name!')
if cmap_trans_fun is None:
cmap_trans_fun = colors.Normalize()
if new_cmap_name is None:
new_cmap_name = cmap_basic.name+'_new'
c_coords_linear = np.linspace(0,1, how_many_levels)
cmap_trans_fun_copy = copy.deepcopy(cmap_trans_fun)
# deppcopy to avoid overwritting the vmin, vmax values
cmap_trans_fun_copy.autoscale([0,1])
c_coords_after = cmap_trans_fun_copy(c_coords_linear)
c_list_after = cmap_basic(c_coords_after)
new_cmap = colors.LinearSegmentedColormap.from_list(new_cmap_name,
c_list_after,
N=how_many_levels)
if ticks_trans_fun is None:
ticks_trans_fun = colors.Normalize()
ticks_trans_fun_copy = copy.deepcopy(ticks_trans_fun)
ticks_trans_fun_copy.vmin = cmap_trans_fun.vmin
ticks_trans_fun_copy.vmax = cmap_trans_fun.vmax
new_ticks_locators = ticks_trans_fun_copy(ticks)
return new_cmap, new_ticks_locators
###########################################
# Prepare some data
# based on https://matplotlib.org/stable/gallery/userdemo/colormap_normalizations.html#sphx-glr-gallery-userdemo-colormap-normalizations-py
N = 100
X, Y = np.mgrid[-3:3:complex(0, N), -2:2:complex(0, N)]
# A low hump with a spike coming out of the top right. Needs to have
# z/colour axis on a log scale so we see both hump and spike. linear
# scale only shows the spike.
Z1 = np.exp(-X**2 - Y**2)
Z1b = 2*np.exp(-10*(X-0.5)**2 - 20*(Y-0.5)**2)
Z2 = np.exp(-(X * 10)**2 - (Y * 10)**2)
Z = Z1 + 50 * Z2+ Z1b
Z = Z*10
# data prepared!
###########################################
cbar_ticks = [0,1,10,50,100,200,300,400]
# prepare basic 'linear' cmap from matplotlib avaliable cmap
cmap = 'inferno'
lin_cmap = cm.get_cmap(cmap)
############################################
# prepare nonlinear norm (good for data viz)
gamma = 0.1
nonlin_norm = colors.PowerNorm(gamma=gamma)
nonlin_norm.autoscale(Z)
# prepare nonlinear norm b (better for colorbar)
gamma_b = 0.40
nonlin_norm_b = colors.PowerNorm(gamma=gamma_b)
nonlin_norm.autoscale(Z)
################# PLOT ####################
# create 4 plots in 2x2 grid
fig, axs = plt.subplots(nrows = 2,
ncols = 2,
figsize=(6,5.5),
dpi=108,
squeeze = False)
fs = 8 # the same plot title fontsize
LLL_ax = axs[0,0]
LNL_ax = axs[0,1]
LNN_ax = axs[1,0]
LNNb_ax = axs[1,1]
#------------- Top left plot --------------
LLL_ax.set_title('linear cmap\nlinear norm\nlinear cbar',
fontsize = fs)
LLL_pcm = LLL_ax.pcolormesh(X, Y, Z,
cmap = lin_cmap,
shading='auto')
# colorbar takes LLL_pcm object to figure out colormap and scale
fig.colorbar(LLL_pcm,
ax=LLL_ax,
extend='both',
ticks=cbar_ticks)
#------------- Top right plot -------------
# an easy way of obtaining good color-scaling
# the colorbar shows cmap in a linear way
# the cbar yticks are nonlinearly scaled but
# they are overlapping
LNL_ax.set_title('linear cmap\nnonlinear norm\nlinear cbar cmap (nonlinear ticks)',
fontsize = fs)
nonlin_norm.autoscale(Z)
LNL_pcm = LNL_ax.pcolormesh(X, Y, Z,
cmap = lin_cmap,
norm = nonlin_norm,
shading='auto')
fig.colorbar(LNL_pcm,
ax=LNL_ax,
extend='both',
ticks=cbar_ticks)
#------------- Bottom left plot -----------
# the colorbar cmap is nonlinear
# the cbar yticks are linearly scaled but
# the overall effect is not good
# the cbar yticks are overlapping again
LNN_ax.set_title('linear cmap\nnonlinear norm\nnonlinear cbar cmap (linear ticks)',
fontsize = fs)
LNN_pcm = LNN_ax.pcolormesh(X, Y, Z,
cmap = lin_cmap,
norm = nonlin_norm,
shading='auto')
# create new, nonlinear cmap
nonlin_cmap, new_ticks_coords = transform_cmap(cmap_basic = lin_cmap ,
cmap_trans_fun = nonlin_norm,
how_many_levels = 256,
ticks = cbar_ticks,
ticks_trans_fun = None,
new_cmap_name = 'nonlinear_cmap')
# create object based on new cmap for colorbar
scalar_mappable = cm.ScalarMappable(cmap=nonlin_cmap)
LNN_cbar = fig.colorbar(scalar_mappable,
ax=LNN_ax,
extend='both',
ticks=new_ticks_coords)
# ticks are in correct places but they are normalized to [0,1] interval
# we need to overwrite them with desired labels
LNN_cbar.ax.set_yticklabels(cbar_ticks)
#------------- Bottom right plot ----------
# the colorbar shows cmap in a nonlinear way
# this is different nonlinear scaling than before (nonlin_norm_b)
# the cbar yticks are also nonlinearly scaled
# this is A GOOD LOOKING PLOT
LNNb_ax.set_title('linear cmap\nnonlinear norm\n2nd nonlinear cbar cmap (nonlinear ticks)',
fontsize = fs)
LNNb_pcm = LNNb_ax.pcolormesh(X, Y, Z,
cmap = lin_cmap,
norm = nonlin_norm,
shading='auto')
# this time as the cbar cmap is with different norm than data cmap
# we also need to recalculate positions of the cbar ticks using second norm
nonlin_cmap_b, new_ticks_coords_b = transform_cmap(cmap_basic = lin_cmap ,
cmap_trans_fun = nonlin_norm_b,
how_many_levels = 256,
ticks = cbar_ticks,
ticks_trans_fun = nonlin_norm_b,
new_cmap_name = 'nonlinear_cmap_v2')
scalar_mappable_b = cm.ScalarMappable(cmap=nonlin_cmap_b)
LNNb_cbar = fig.colorbar(scalar_mappable_b,
ax=LNNb_ax,
extend='both',
ticks=new_ticks_coords_b)
LNNb_cbar.ax.set_yticklabels(cbar_ticks)
#------------------------------
fig.tight_layout()
plt.show()
I was using this answer as a base:
Uniform tick labels for non-linear colorbar in Matplotlib
These answers may be useful but were looking too complicated:
Arbirtrary non-linear colorbar using Matplotlib
nonlinear colormap, matplotlib
I have a feeling that wiser usage of norm parameter in pcolor and perhaps in cbar should give me the desired result. Unfortunately, I was not able to obtain it in this way.
