plot_tfr_stub.py

def plot_tfr(tfr, vmin, vmax, axes=None, colorbar=False, title=None,
             set_chan_title=True, show=True):
    from mne.viz.utils import center_cmap
    from matplotlib import pyplot as plt
    cmap = center_cmap(plt.cm.RdBu_r, vmin, vmax)  # zero maps to white
    n_chans = len(tfr.ch_names)
    if axes is None:
        if colorbar:
            w_ratios = list(np.repeat(10, n_chans)) + [1]
            fig, axes = plt.subplots(1, n_chans + 1, figsize=(12, 4),
                                     gridspec_kw={"width_ratios": w_ratios})
        else:
            fig, axes = plt.subplots(1, n_chans, figsize=(12, 4))
    else:
        fig = axes[0].get_figure()

    if colorbar:
        assert len(axes) == n_chans + 1
    else:
        assert len(axes) == n_chans

    for i, ch in enumerate(tfr.ch_names):  # for each channel
        ax = axes[i]
        tfr.plot([ch], vmin=vmin, vmax=vmax, cmap=(cmap, False),
                 axes=ax, colorbar=False, show=False, title='')
        if set_chan_title:
            ax.set_title(ch, fontsize=10)
        if not ax.is_first_col():
            ax.set_ylabel("")
            ax.set_yticklabels("")
    if colorbar:
        fig.colorbar(axes[0].images[-1], cax=axes[-1])
    if title is not None:
        fig.suptitle(title)
    if show:
        plt.show()
    return fig