Broken colorbar in matplotlib

broken_colorbar.png

broken_colorbar.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.gridspec as gs
import matplotlib.patches as patches

# Data points
x = y = np.linspace(-4*np.pi,4*np.pi,2000)
X,Y = np.meshgrid(x,y)
Z1 = 10 + 10*(np.sin(X-np.pi/2) + np.cos(Y))
Z2 = -5 + 2*np.cos(Y)

# Select peaks of the first functions, the rest is filled by the second function
mask = Z1<10
Z1 = np.ma.MaskedArray(Z1, mask)
Z2 = np.ma.MaskedArray(Z2, np.logical_not(mask))

# Create figure and subplots with gridspec
fig = plt.figure()
g = gs.GridSpec(2,2, width_ratios=(1,0.02), hspace=0)
ax = fig.add_subplot(g[:,0])
ax.set_aspect('equal')
ax.set_xticks(np.pi*np.arange(-4,5,1), [(r'%d$\pi$'%n).replace('1','').replace('0$\pi$','0') for n in np.arange(-4,5,1)])
ax.set_yticks(np.pi*np.arange(-4,5,1), [(r'%d$\pi$'%n).replace('1','').replace('0$\pi$','0') for n in np.arange(-4,5,1)])

cmap1 = plt.get_cmap('Reds')
cmap2 = plt.get_cmap('Blues_r') #_r reverses the map so whites are at the end

# Remove the white part of the colormaps so that it doesn't appear like the values coincide
# By unutbu https://stackoverflow.com/a/18926541
def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
    new_cmap = colors.LinearSegmentedColormap.from_list(
        'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
        cmap(np.linspace(minval, maxval, n)))
    return new_cmap

cmap1 = truncate_colormap(cmap1, minval=0.2)
cmap2 = truncate_colormap(cmap2, maxval=0.8)

im1 = plt.contourf(X,Y,Z1, cmap=cmap1, levels=np.linspace(10,30,50))
im2 = plt.contourf(X,Y,Z2, cmap=cmap2, levels=np.linspace(-7,-3,50)) # _r is to reverse the colormap

# Bottom colorbar
cbax1 = fig.add_subplot(g[1,1])
cbar1 = fig.colorbar(im2, cax=cbax1, ticks=np.arange(-7,-2,1))

# Top colorbar
cbax2 = fig.add_subplot(g[0,1])
cbar2 = fig.colorbar(im1, cax=cbax2, ticks=np.arange(10,35,5))

# Adjust colorbar positions and add a spacer in between
_,b1,w,h = cbax1.get_position().bounds  # left,bottom,width,height
_,b2,w,h = cbax2.get_position().bounds
R = ax.get_position().max[0] # right coordinate of main plot

hs = 0.2*h # height of the spacer
cbax1.set_position([1.05*R, b1,      w, h-hs/2])
cbax2.set_position([1.05*R, b2+hs/2, w, h-hs/2])

# Spacer is a white rectangle patch with some dots
sax = fig.add_axes([1.05*R, b1+h-hs/2, w, hs])
sax.add_patch(patches.Rectangle((-1,-1),2,6, fill=False, edgecolor=None))
sax.plot([0,0,0],[1,2,3],'k.',ms=2)
sax.set_xlim([-0.5,0.5])
sax.set_ylim([0,4])
sax.xaxis.set_visible(False)
sax.yaxis.set_visible(False)

# Remove some spines (box edges lines)
# spine.set_visible(False) does not seem to work on the colorbar axes, so the easiest
# is to disable the outlines for everything, then redraw a box around
cbar1.outline.set_visible(False)
cbar2.outline.set_visible(False)
for spine in sax.spines.values():
    spine.set_visible(False)

bax = fig.add_axes([1.05*R,b1,w,2*h])
bax.xaxis.set_visible(False)
bax.yaxis.set_visible(False)
bax.set_facecolor("none")

# Title
fig.suptitle(r"$z = 10+10[\sin(x-\pi/2)+\cos(y)]$ if $z\geq 10$,"
              "\n"
             r"$-5+2\cos(y)$ otherwise")

# Save 
fig.savefig("broken_colorbar.png", dpi=100, bbox_inches="tight")