Make matplotlib.pyplot color bar span two rows alongside waveform and specgram

Question

Using matplotlib.pyplot, I have two plots. One is a waveform of an audio file. The second is a spectrogram of the same audio. I want the wave form to be directly above the spectrogram (same x-axis, and aligned together). I also want a colorbar for the spectrogram.

Problem - when I put the colorbar in, it attaches to the spectrogram row and the waveform extends over the colorbar (i.e. is no longer time-aligned with the spectrogram and is wider than the spectrogram).

I am close to the solution, I think, but I just can't quite figure out what I'm doing wrong or what to change to get it working the way I want. Hope someone can point me in the right direction!

Using the following python code (I made the code as MWE as possible):

import matplotlib
matplotlib.use("TkAgg")
from scipy.io import wavfile
from matplotlib import mlab
from matplotlib import pyplot as plt
import numpy as np
from numpy.lib import stride_tricks

samplerate, data = wavfile.read('FILENAME.wav')

times = np.arange(len(data))/float(samplerate)

plt.close("all")

####
#Waveform
####
fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(13.6, 7.68))

plt.subplot(211)
plt.plot(times, data, color='k') 

plt.xlabel('time (s)')
plt.xlim(times[0], times[-1])

max_amp = max(abs(np.amin(data)), abs(np.amax(data)))
min_amp = (max_amp * -1) - abs(np.amin(data) - np.amax(data))/50
max_amp = max_amp + abs(np.amin(data) - np.amax(data))/50

plt.ylim(min_amp, max_amp)

ax = plt.gca()
ax.set_yticks(np.array([min_amp, min_amp/2, 0, max_amp/2, max_amp]))

ax.spines['bottom'].set_position('center')
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')

ax.xaxis.set_ticks_position('none')
ax.yaxis.set_ticks_position('none')

ax.xaxis.set_tick_params(pad=115)

####
#Spectrogram
####

Fs = 5000*2.#10000.
NFFT = min(512, len(data))
noverlap = NFFT / 2
pad_to = NFFT * 16
dynamicRange = 27.5
vmin = 20*np.log10(np.max(data)) - dynamicRange

cmap = plt.get_cmap('inferno')

plt.subplot(212)
Pxx, freqs, times, cax = plt.specgram(data, NFFT=NFFT, Fs=samplerate, noverlap=noverlap, mode='magnitude', scale='dB', vmin=vmin, pad_to=pad_to, cmap=cmap)

axes_spec = plt.gca()
axes_spec.set_xlim([0., max(times)])
axes_spec.set_ylim([0, 5000])

plt.xlabel("Time (s)")
plt.ylabel("Frequency (hz)")

plt.colorbar(cax, label='(dB)').ax.yaxis.set_label_position('left')


plt.tight_layout()
plt.show()

I can get the following plot:

Making these slight modifications below, I can get the plot to look almost how I want. The problem is, it creates a blank figure next to the colorbar. This version, minus the blank figure, is what I am trying to create.

#Replace this for waveform
plt.subplot(221)
#Replace this for spectrogram
plt.subplot(223)
#Add this before colorbar
plt.subplot(122)

New version of plot:

EDIT: There is another possibility that I am also OK with (or perhaps both, for good measure!)

Answer 1

Here is an example of colorbar based on one of the answers in matplotlib-2-subplots-1-colorbar. The parameter pad in fig.colorbar is used to specify the space between the plots and the colorbar, and aspect is used to specify the aspect ratio between the height and width of the colorbar. Specgram outputs the image as the 4th output parameter, so I'm using that for the colorbar.

fig,axs = matplotlib.pyplot.subplots(ncols=1, nrows=2 )
N=1000; fs=10e3
x = np.sin(np.arange(N))+np.random.random(N)
spectrum, freqs, t, im = axs[1].specgram(x,Fs=fs,
                    cmap=matplotlib.cm.inferno,noverlap=255)
axs[0].plot(np.arange(0,N)/fs,x,'-');
axs[0].set_xlim(t[0],t[-1]);axs[1].set_xlim(t[0],t[-1])
axcb = fig.colorbar(im, ax=axs.ravel().tolist(), pad=0.04, aspect = 30)

It is important to notice that when fig.colorbar function is called using the ax parameter, the original plots will be resized to make room for the colorbar. If it is only applied to one of the plots, only that axis will be resized. Se below:

fig,axs = matplotlib.pyplot.subplots(ncols=1, nrows=2 )
N=1000; fs=10e3
x = np.sin(np.arange(N))+np.random.random(N)
spectrum, freqs, t, im = axs[1].specgram(x,Fs=fs,
                    cmap=matplotlib.cm.inferno,noverlap=255)
axs[0].plot(np.arange(0,N)/fs,x,'-')
axs[0].set_xlim(t[0],t[-1]);axs[1].set_xlim(t[0],t[-1])
axcb = fig.colorbar(im, ax=axs[1], pad=0.04, aspect = 30)

Below it is shown a way of controlling the resizing of your original axes in order to make room for a colorbar using fig.colorbar with the cax parameter that will not resize further your original plots. This approach requires to manually make some room for your colorbar specifying the right parameter inside the function fig.subplots_adjust :

fig,axs = matplotlib.pyplot.subplots(ncols=1, nrows=2 )
N=1000; fs=10e3
x = np.sin(np.arange(N))+np.random.random(N)
spectrum, freqs, t, im = axs[1].specgram(x,Fs=fs,
                    cmap=matplotlib.cm.inferno,noverlap=255)
axs[0].plot(np.arange(0,N)/fs,x,'-')
axs[0].set_xlim(t[0],t[-1]);axs[1].set_xlim(t[0],t[-1])
fig.subplots_adjust(right=0.85)  # making some room for cbar
# getting the lower left (x0,y0) and upper right (x1,y1) corners:
[[x10,y10],[x11,y11]] = axs[1].get_position().get_points()
pad = 0.01; width = 0.02
cbar_ax = fig.add_axes([x11+pad, y10, width, y11-y10])
axcb = fig.colorbar(im, cax=cbar_ax)

And doing the same to span two rows by reading coordinates of the original two plots:

fig,axs = matplotlib.pyplot.subplots(ncols=1, nrows=2 )
N=1000; fs=10e3
x = np.sin(np.arange(N))+np.random.random(N)
spectrum, freqs, t, im = axs[1].specgram(x,Fs=fs,
                    cmap=matplotlib.cm.inferno,noverlap=255)
axs[0].plot(np.arange(0,N)/fs,x,'-')
axs[0].set_xlim(t[0],t[-1]);axs[1].set_xlim(t[0],t[-1])
fig.subplots_adjust(right=0.85)  # making some room for cbar
# getting the lower left (x0,y0) and upper right (x1,y1) corners:
[[x00,y00],[x01,y01]] = axs[0].get_position().get_points()
[[x10,y10],[x11,y11]] = axs[1].get_position().get_points()
pad = 0.01; width = 0.02
cbar_ax = fig.add_axes([x11+pad, y10, width, y01-y10])
axcb = fig.colorbar(im, cax=cbar_ax)

Answer 2

The best solution I came up with is subplot2grid() function. This requies the use of subplots, which I was not using originally. Following this method, I needed to change everything from using plt (matplotlib.pyplot) to using the axes for the given plot for each .plot() or .specgram() invocation. The relevant changes are included here:

#No rows or columns need to be specified, because this is handled within a the `subplot2grid()` details
fig, axes = plt.subplots(figsize=(13.6, 7.68))

#Setup for waveform
ax1 = plt.subplot2grid((2, 60), (0, 0), rowspan=1,  colspan=56)
####WAVEFORM PLOTTING

#Setup for spectrogram
ax2 = plt.subplot2grid((2, 60), (1, 0), rowspan=1,  colspan=56)
####SPECTROGRAM PLOTTING

#Setup for colorbar
ax3 = plt.subplot2grid((2, 60), (0, 59), rowspan=1, colspan=1)
cbar = plt.colorbar(cax, cax=ax3, ax=ax2)

And a MWE bringing it all together:

import matplotlib as mpl
mpl.use("TkAgg")
from scipy.io import wavfile
from matplotlib import mlab
from matplotlib import pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
from numpy.lib import stride_tricks

samplerate, data = wavfile.read('FILENAME.wav')

times = np.arange(len(data))/float(samplerate)

plt.close("all")

fig, axes = plt.subplots(figsize=(13.6, 7.68))#nrows=2, ncols=2, 

gs = gridspec.GridSpec(2, 60)

####
#Waveform
####
ax1 = plt.subplot2grid((2, 60), (0, 0), rowspan=1,  colspan=56)
ax1.plot(times, data, color='k')

ax1.xaxis.set_ticks_position('none')
ax1.yaxis.set_ticks_position('none')

####
#Spectrogram
####
maxFrequency = 5000
Fs = maxFrequency*2.#10000.
NFFT = min(512, len(data))
noverlap = NFFT / 2
pad_to = NFFT * 16
dynamicRange = 27.5
vmin = 20*np.log10(np.max(data)) - dynamicRange

cmap = plt.get_cmap('inferno')

ax2 = plt.subplot2grid((2, 60), (1, 0), rowspan=1,  colspan=56)
Pxx, freqs, times, cax = ax2.specgram(data, NFFT=NFFT, Fs=samplerate, noverlap=noverlap, mode='magnitude', scale='dB', vmin=vmin, pad_to=pad_to, cmap=cmap)

ax2.set_ylim([0, maxFrequency])
ax2.xaxis.set_ticks_position('none')
ax2.yaxis.set_ticks_position('none')

####
#Colorbar (for spectrogram)
####
ax3 = plt.subplot2grid((2, 60), (1, 59), rowspan=1, colspan=1)
cbar = plt.colorbar(cax, cax=ax3, ax=ax2)
cbar.ax.yaxis.set_tick_params(pad=3, left='off', right='off', labelleft='on', labelright='off')

plt.show()

Here's an example of the output from this MWE:

Best part! You need only change the 0 to 1 and the rowspan to be 1 in this line (i.e. :)

ax3 = plt.subplot2grid((2, 60), (1, 59), rowspan=1, colspan=1)

to make the colorbar span only the height of the spectrogram. Meaning that changing between the two options is incredibly simple. Here's an example of the output from this change:

EDIT: GridSpec actually was unused, and so I edited it out. The only relevant details that I needed involved calling subplot2grid() to set up the subplots.