Calculate adaptive mean amplitude

adaptive_mean.py

import os.path as op
import numpy as np
import mne
from mne.datasets import sample

# Import sample data
path = op.join(sample.data_path(), 'MEG', 'sample')
raw = mne.io.read_raw_fif(op.join(path, 'sample_audvis_filt-0-40_raw.fif'))
events = mne.read_events(op.join(path, 'sample_audvis_filt-0-40_raw-eve.fif'))
epochs = mne.Epochs(raw, events, [1, 2], tmin=-0.2, tmax=1.0, preload=True)

# Define search window for the peak
search_window = (0.05, 0.10)  # In seconds

# Compute the corresponding samples
ind_start, ind_stop = np.searchsorted(epochs.times, search_window)

# Get the data in the search window as a numpy array
data_to_search = epochs.get_data()[:, :, ind_start:ind_stop]

# For each epoch and each channel, find the maximum value in the search window
peaks = data_to_search.argmax(axis=2)

# The indices are currently based on the data_to_search array. Add the index
# of the start of the search window to the peaks to obtain indexes that can be
# used with the original epochs data.
peaks += ind_start

# Define the window around the peak for which we want to compute the mean,
# where t=0 is the peak itself.
peak_window = (-0.1, 0.1)  # In seconds

# Convert the peak window into an array of sample indices
peak_window = np.arange(int(peak_window[0] * epochs.info['sfreq']),
                        int(peak_window[1] * epochs.info['sfreq']))

# Use numpy broadcasting to add the peaks to the peak_window. This gives us for
# each epoch and channel, the corresponding time window over which to compute
# the mean.
time_windows = peaks[:, :, np.newaxis] + peak_window[np.newaxis, np.newaxis, :]

# Use advanced integer indexing to get the corresponding data from the epochs
n_epochs, n_channels, n_samples = epochs.get_data().shape
peak_data = epochs.get_data()[
    np.arange(n_epochs)[:, np.newaxis, np.newaxis],  # All epochs
    np.arange(n_channels)[np.newaxis, :, np.newaxis],  # All channels
    time_windows  # Only the requested time windows
]

# Finally, compute the mean across time
means = peak_data.mean(axis=2)