Script to plot Power Spectral Density (PSD) from raw data.

psd_plot.py

#!/usr/bin/env python

# import the necessary packages
import numpy as np
import mne
import matplotlib.pyplot as pl
import sys

# turn interactive plotting on
pl.ion()

# read the raw file as an argument to the script
# e.g. from the terminal
# python psd_plot.py /home/kate/<raw file name-raw.fif>
# or from the ipython console
# run psd_plot.py /home/kate/<raw file name-raw.fif>
raw_fname = sys.argv[1]

# reading the raw file
raw = mne.io.Raw(raw_fname, preload=True)
# selecting the channels of interest
# here we choose only the meg channels
picks = mne.pick_types(raw.info, meg=True, eeg=False, eog=False,
                       ecg=False, exclude=[])

# select bad picks
bads = raw.info['bads']
bad_picks = [raw.ch_names.index(l) for l in bads]

# you can use the below code
# 1. to plot average power spectral density over all the channels
# raw.plot_psd(fmin=2., fmax=200., picks=None, n_jobs=1)

# 2. or plot average psd for just one channel
# mypick = [0]
# raw.plot_psd(fmin=2., fmax=200., picks=mypick, n_jobs=1)

# compute power spectral densities for all channels, but do not plot
# it will be stored as an array under psds
psds, freqs = mne.time_frequency.psd_welch(raw, picks=picks,
                                           fmin=2., fmax=200.)

# plot all the psds in the same figure
# the channels already marked as bad will be plotted in red
pl.figure('psd butterfly')
for k in picks:
    if k in bad_picks:
        pl.plot(freqs, 10 * np.log10(psds[k]), color='r')
    else:
        pl.plot(freqs, 10 * np.log10(psds[k]), color='b')

# save the figure, change the filename if necessary
pl.savefig('psd_butterfly.png')