Keep the linear colors of a colormap while using non-linear levels in plt.contourf

Question

In the following example, I would like to manually set the levels of a contourf plot to the quantiles of some variable (in order to enhance the vizualization).

But giving the levels also affects the scaling of the colors of the colormap (see how the red is extended -> the visual impression is not affected): how can I keep them linear?

In other words, how can I keep the colors of the colorbar just as in the top subplot, but corresponding to the values as in the bottom subplot?

import matplotlib.pyplot as plt
import numpy as np

xvec = np.linspace(-10.,10.,100)
x,y = np.meshgrid(xvec, xvec)
z = -(x**2 + y**2)**2

fig, (ax0, ax1) = plt.subplots(2)
p0 = ax0.contourf(x, y, z, 100)
fig.colorbar(p0, ax=ax0)
p1 = ax1.contourf(x, y, z, 100, levels=np.percentile(z, np.linspace(0,100,101)))
fig.colorbar(p1, ax=ax1)

plt.show()

Did I miss some easy solution (I bet yes), or should I try to do some matplotlib.colors.LinearSegmentedColormap manipulation... ?!

Answer 1

You could use a matplotlib.colors.BoundaryNorm to specify the levels to use for the colormapping.

import matplotlib.pyplot as plt
import matplotlib.colors 
import numpy as np

xvec = np.linspace(-10.,10.,100)
x,y = np.meshgrid(xvec, xvec)
z = -(x**2 + y**2)**2

fig, (ax0, ax1) = plt.subplots(2)
p0 = ax0.contourf(x, y, z, 100)
fig.colorbar(p0, ax=ax0)

levels = np.percentile(z, np.linspace(0,100,101))
norm = matplotlib.colors.BoundaryNorm(levels,256)
p1 = ax1.contourf(x, y, z, 100, levels=levels, norm=norm)
fig.colorbar(p1, ax=ax1)

plt.show()