peak to peak annotations

ptp_annotations.py

import mne

import numpy as np


def _append_annotation(raw, new_events, onset, duration,extend_duration=0.1):
    onset_sec = (onset - raw.first_samp) / raw.info['sfreq']
    onset_sec -= extend_duration
    new_events['onset'].append(onset_sec)
    new_events['duration'].append(duration + 2 * extend_duration)

def get_ptp_annotations(raw, ptp_thresh=0.5e-6, duration=0.05, extend_duration=0.1):

    # We want to get annotations for high-amplitude periods.  We can accomplish this by
    # creating epochs of very tiny durations and checking whether the epoch is
    # greater than some threshold. We then want to extend that high-amplitude period marked
    # by a few seconds on both sides of the annotations

    # Create fixed length events and create epochs in small duratins
    events_fixed = mne.make_fixed_length_events(raw,
         id=1, duration=duration, first_samp=True, overlap=0.0)
    epochs = mne.Epochs(raw, events_fixed, tmin=0, tmax=duration,
                        picks='meg', baseline=None)
    annotations = None

    # n_epochs x n_channels
    # Check peach to peak amplitude and get an index/events for when that epoch
    # is below the designated threshold

    epochs_ptp = np.ptp(epochs.get_data(), axis=-1)
    ptp_idx = np.where(np.min(epochs_ptp, axis=-1) > ptp_thresh)[0]
    ptp_events = events_fixed[ptp_idx]

    # merge events
    this_duration = duration

    # Check that this function has detected bad periods of data
    if ptp_events.size != 0:

        # Create onset from first event and empty dictionary
        onset = ptp_events[0][0]
        new_events = dict(onset=list(), duration=list())

        # iterate through the bad events detected
        for idx, event in enumerate(ptp_events[:-1]):
            next_samp = ptp_events[idx + 1][0]

            if onset + this_duration * raw.info['sfreq'] >= next_samp:
                this_duration += duration

                if idx == len(ptp_events) - 2:
                    _append_annotation(raw, new_events, onset, this_duration,
                                       extend_duration)
            else:
                _append_annotation(raw, new_events, onset, this_duration,
                                   extend_duration)

                # reset duration and onset for next annotation
                this_duration = duration
                onset = next_samp

        # Double check that when epoching, MNE rejects BAD_PTP
        # Create annotation object to set on raw data
        annotations = mne.Annotations(onset=new_events['onset'],
                                      duration=new_events['duration'],
                                      description='BAD_PTP')

    return annotations