Skip to content

Instantly share code, notes, and snippets.

@drammock
Last active April 24, 2020 12:06
Show Gist options
  • Save drammock/cbf67cbfd65736b7ef028581e2240d55 to your computer and use it in GitHub Desktop.
Save drammock/cbf67cbfd65736b7ef028581e2240d55 to your computer and use it in GitHub Desktop.
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: Daniel McCloy
Load SSVEP epochs and plot PSDs, phases, etc.
data at https://dan.mccloy.info/data/prek_1964-pre_camp-pskt-epo.fif
"""
import os
import numpy as np
from scipy.ndimage import convolve1d
import matplotlib.pyplot as plt
import mne
from mne.time_frequency.multitaper import (_compute_mt_params, _mt_spectra,
_psd_from_mt)
# config
in_dir = '/home/drmccloy/Desktop'
fig_fname = os.path.join(in_dir, 'phases.pdf')
s = 'prek_1964'
timepoint = 'pre'
bandwidth = 0.1
subdiv = ''
# load epochs
fname = f'{s}-{timepoint}_camp-pskt{subdiv}-epo.fif'
epochs = mne.read_epochs(os.path.join(in_dir, fname), proj=True)
# epochs.pick_types(meg=True) # drop EOG, ECG channels
evoked = epochs.average()
n_times = len(evoked.times)
# do multitaper estimation
sfreq = evoked.info['sfreq']
mt_kwargs = dict(n_times=n_times, sfreq=sfreq, bandwidth=bandwidth,
low_bias=True, adaptive=False)
dpss, eigvals, adaptive = _compute_mt_params(**mt_kwargs)
n_fft = n_times # _mt_spectra defaults to wrong axis
mt_spectra, freqs = _mt_spectra(evoked.data, dpss, sfreq, n_fft=n_fft)
magnitudes = np.abs(mt_spectra)
phases = np.angle(mt_spectra)
# compute the sensor-space PSD
sensor_weights = np.sqrt(eigvals)[np.newaxis, :, np.newaxis]
sensor_psd = _psd_from_mt(mt_spectra, sensor_weights)
# divide each bin by its two neighbors on each side
snr_psd = sensor_psd / convolve1d(sensor_psd, mode='constant',
weights=[0.25, 0.25, 0, 0.25, 0.25])
these_freqs = (2, 3, 4, 5, 6, 11, 12)
n_freqs = len(these_freqs)
fig = plt.figure(figsize=(15, 4 * n_freqs))
n_columns = 4
for ix, freq in enumerate(these_freqs):
# find the sensor that best captures {freq} Hz activity
bin_idx = np.argmin(np.abs(freqs - freq))
best_sensor = np.argmax(snr_psd[..., bin_idx], axis=0)
best_sensor_name = evoked.ch_names[bin_idx]
ax = plt.subplot(n_freqs, n_columns, n_columns * ix + 1)
ax.plot(freqs, snr_psd[best_sensor].T, linewidth=1, color='k')
ax.set(xticks=np.arange(0, 25, 2), xlabel='freq (Hz)',
ylabel=f'{freq} Hz\n"SNR" (a.u.)', ylim=(0, 40))
title = f'best {freq} Hz SNR: {best_sensor_name}'
if not ix:
title = f'PSD divided by sum of 2 bins on each side\n(channel w/ {title})' # noqa E501
ax.set(title=title)
# phase vs magnitude
these_phases = np.squeeze(phases[..., bin_idx])
these_magnitudes = np.squeeze(magnitudes[..., bin_idx])
ax = plt.subplot(len(these_freqs), n_columns, n_columns * ix + 2)
ax.plot(these_phases, these_magnitudes, '.', alpha=0.5)
ax.set(xlim=(-np.pi, np.pi), xlabel='phase', ylabel='magnitude',
ylim=(0, magnitudes.max()))
if not ix:
ax.set(title='phase vs magnitude')
# phase vs magnitude (polar)
ax = plt.subplot(n_freqs, n_columns, n_columns * ix + 3, polar=True)
ax.plot(these_phases, these_magnitudes, '.', alpha=0.5)
ax.set(ylim=(0, magnitudes.max()))
if not ix:
ax.set(title='phase vs magnitude\n')
# phase histogram
ax = plt.subplot(n_freqs, n_columns, n_columns * ix + 4, polar=True)
ax.hist(these_phases, bins=72)
ax.set(ylim=(0, 20))
if not ix:
ax.set(title='phase histogram (5° bins)\n')
fig.tight_layout()
fig.subplots_adjust(top=0.95, left=0.05, hspace=0.5)
fig.savefig(fig_fname)
plt.close('all')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment