python - Colorbar limits are not respecting set vmin/vmax in plt.contourf. How can I more explicitly set the colorbar limits?

Question

Getting a strange result when trying to adjust the data range when plotting using contourf

import matplotlib
import numpy as np
import matplotlib.cm as cm
import matplotlib.mlab as mlab
import matplotlib.pyplot as plt

delta = 0.025
x = np.arange(-3.0, 3.0, delta)
y = np.arange(-2.0, 2.0, delta)
X, Y = np.meshgrid(x, y)
Z1 = mlab.bivariate_normal(X, Y, 1.0, 1.0, 0.0, 0.0)
Z2 = mlab.bivariate_normal(X, Y, 1.5, 0.5, 1, 1)
# difference of Gaussians
Z = 10.0 * (Z2 - Z1)

plt.figure()
CS = plt.contourf(X, Y, Z, vmin = 0, vmax = 3)
plt.title('Simplest default with labels')
plt.colorbar()

plt.show()

Results in this for me:

It's like the colors match the vmin/vmax I set, but the number range displayed on the colorbar remains what it would be without setting vmin/vmax.

In this case, I would like the end result to have a colorbar that ranges from 0 to 3.

Answer 1

First of all, the response, marked as answer, is erroneous (see my comments above), but helped me to come up with two other solutions.

As JulianBauer pointed out in a comment below, the function mlab.bivariate_normal used by the OP is not available any more. To provide functional code that produces output that can be compared with the other answers I am calling the following function, with the definition of bivariate_normal copied from the matplotlib repository:

def myfunction():

    def bivariate_normal(X, Y, sigmax=1.0, sigmay=1.0, mux=0.0, muy=0.0, sigmaxy=0.0):
        """copied from here: https://github.com/matplotlib/matplotlib/blob/81e8154dbba54ac1607b21b22984cabf7a6598fa/lib/matplotlib/mlab.py#L1866"""
        Xmu = X-mux
        Ymu = Y-muy
        rho = sigmaxy/(sigmax*sigmay)
        z = Xmu**2/sigmax**2 + Ymu**2/sigmay**2 - 2*rho*Xmu*Ymu/(sigmax*sigmay)
        denom = 2*np.pi*sigmax*sigmay*np.sqrt(1-rho**2)
        return np.exp(-z/(2*(1-rho**2))) / denom

    delta = 0.025
    x = np.arange(-3.0, 3.0, delta)
    y = np.arange(-2.0, 2.0, delta)
    X, Y = np.meshgrid(x, y)
    Z1 = bivariate_normal(X, Y, 1.0, 1.0, 0.0, 0.0)
    Z2 = bivariate_normal(X, Y, 1.5, 0.5, 1, 1)
    Z = 10.0 * (Z2 - Z1)
    return X,Y,Z

1. A simple and straight forward solution

Make use of the extend command while providing custom levels:

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

X,Y,Z = myfunction()

plt.figure()
plt.title('Simplest default with labels')
levels = np.linspace(0.0, 3.0, 7)
CS = plt.contourf(X, Y, Z, levels=levels, cmap=cm.coolwarm, extend='min')

colorbar = plt.colorbar(CS)

plt.show()

2. A more complicated solution

is provided in the answer above, though it needs to be adapted to specific cases and one can easily end up with a colorbar whose levels differs from those in the actual plot. I find this dangerous, so I attempted to wrap it up in a function that can safely be called in any context:

def clippedcolorbar(CS, **kwargs):
    from matplotlib.cm import ScalarMappable
    from numpy import arange, floor, ceil
    fig = CS.ax.get_figure()
    vmin = CS.get_clim()[0]
    vmax = CS.get_clim()[1]
    m = ScalarMappable(cmap=CS.get_cmap())
    m.set_array(CS.get_array())
    m.set_clim(CS.get_clim())
    step = CS.levels[1] - CS.levels[0]
    cliplower = CS.zmin<vmin
    clipupper = CS.zmax>vmax
    noextend = 'extend' in kwargs.keys() and kwargs['extend']=='neither'
    # set the colorbar boundaries
    boundaries = arange((floor(vmin/step)-1+1*(cliplower and noextend))*step, (ceil(vmax/step)+1-1*(clipupper and noextend))*step, step)
    kwargs['boundaries'] = boundaries
    # if the z-values are outside the colorbar range, add extend marker(s)
    # This behavior can be disabled by providing extend='neither' to the function call
    if not('extend' in kwargs.keys()) or kwargs['extend'] in ['min','max']:
        extend_min = cliplower or ( 'extend' in kwargs.keys() and kwargs['extend']=='min' )
        extend_max = clipupper or ( 'extend' in kwargs.keys() and kwargs['extend']=='max' )
        if extend_min and extend_max:
            kwargs['extend'] = 'both'
        elif extend_min:
            kwargs['extend'] = 'min'
        elif extend_max:
            kwargs['extend'] = 'max'
    return fig.colorbar(m, **kwargs)

The main commands in the function correspond to what kilojoules proposes in his/her answer, but more lines are required to avoid all the explicit and potentially erroneous assignments by extracting all information from the contourf object.

Usage:

The OP asks for levels from 0 to 3. The darkest blue represents values below 0, so I find an extend-marker useful.

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

X,Y,Z = myfunction()

plt.figure()
plt.title('Simplest default with labels')
CS = plt.contourf(X, Y, Z, levels=6, vmin=0.0, vmax=3.0, cmap=cm.coolwarm)

colorbar = clippedcolorbar(CS)

plt.show()

The extend marker can be disabled by calling clippedcolorbar(CS, extend='neither') instead of clippedcolorbar(CS).