Skip to content

Instantly share code, notes, and snippets.

@bloyl
Created July 16, 2020 01:43
Show Gist options
  • Save bloyl/a08a64899f2feffda5271890a5c6a677 to your computer and use it in GitHub Desktop.
Save bloyl/a08a64899f2feffda5271890a5c6a677 to your computer and use it in GitHub Desktop.
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment