phase_mwe.py

#!/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')