Read Raw EDF to MNE with different sampling frequencies

read_raw_edf.py

from pyedflib import highlevel
import numpy as np
import mne


def resample(data, o_sfreq, t_sfreq):
    """
    resample a signal using MNE resample functions
    This automatically is optimized for EEG applying filters etc

    Parameters
    ----------
    data : np.ndarray
        a 1D data array.
    o_sfreq : float or int
        the original sampling frequency.
    t_sfreq : float or int
        the target sampling frequency.

    Returns
    -------
    np.ndarray
        the resampled array

    """

    if o_sfreq==t_sfreq: return data
    raw = np.atleast_2d(data)
    ch_names=['ch{}'.format(i) for i in range(len(raw))]
    info = mne.create_info(ch_names=ch_names, sfreq=o_sfreq, ch_types=['eeg'])
    raw_mne = mne.io.RawArray(raw, info, verbose='WARNING')
    resampled = raw_mne.resample(t_sfreq, n_jobs=1, verbose='WARNING')
    new_raw = resampled.get_data().squeeze()
    return new_raw.astype(raw.dtype, copy=False)


def load_edf(edf_file, channels=None, resample_to=max):
    """
    convenience function to load an EDF+ with different sampling frequencies
    without MNE. Workaround for https://github.com/mne-tools/mne-python/issues/10635
    
    Parameters
    ----------
    edf_file : str
        edf file name.
    channels : list of str
        list of channel names to load. If None, all channels are loaded
    resample_to : callable
        function that determines the sampling frequency that will be used as 
        a target. I.e. min or max or lambda x: np.mean(x)

    Returns
    -------
    raw : mne.io.BaseRaw
        raw file containing the signal.
    """
    assert callable(resample_to), 'resample_to must be a function (e.g. min or max)'
    if not isinstance(channels, list):
        channels = [channels]
        
    sigs, sigheads, header = highlevel.read_edf(edf_file, 
                                               ch_names=channels)
   
    sfreqs = [s['sample_rate'] for s in sigheads]
    labels = [s['label'] for s in sigheads]
    
    if len(set(sfreqs))>1:
        sigs = [resample(x, sfreqs[i], resample_to(sfreqs)) for i, x in enumerate(sigs)]

    info = mne.create_info(labels, resample_to(sfreqs), ch_types='eeg')
    raw = mne.io.RawArray(sigs, info)
    return raw