<commit_before>"""Xdawn implementation."""
# Authors: <NAME> <<EMAIL>>
#
# License: BSD (3-clause)

import copy as cp

import numpy as np
from scipy import linalg

from ..io.base import _BaseRaw
from ..epochs import _BaseEpochs
from .. import Covariance, EvokedArray, Evoked, EpochsArray
from ..io.pick import pick_types
from .ica import _get_fast_dot
from ..utils import logger
from ..cov import _regularized_covariance
from ..channels.channels import ContainsMixin


def _least_square_evoked(data, events, event_id, tmin, tmax, sfreq):
    """Least square estimation of evoked response from data.

    Parameters
    ----------
    data : ndarray, shape (n_channels, n_times)
        The data to estimates evoked
    events : ndarray, shape (n_events, 3)
        The events typically returned by the read_events function.
        If some events don't match the events of interest as specified
        by event_id, they will be ignored.
    event_id : dict
        The id of the events to consider
    tmin : float
        Start time before event.
    tmax : float
        End time after event.
    sfreq : float
        Sampling frequency.

    Returns
    -------
    evokeds_data : dict of ndarray
        A dict of evoked data for each event type in event_id.
    toeplitz : dict of ndarray
        A dict of toeplitz matrix for each event type in event_id.
    """
    nmin = int(tmin * sfreq)
    nmax = int(tmax * sfreq)

    window = nmax - nmin
    n_samples = data.shape[1]
    toeplitz_mat = dict()
    full_toep = list()
    for eid in event_id:
        # select events by type
        ix_ev = events[:, -1] == event_id[eid]

        # build toeplitz matrix
        trig = np.zeros((n_samples, 1))
        ix_trig = (events[ix_ev, 0]) + nmin
        trig[ix_trig] = 1
        toep_mat = linalg.toeplitz(trig[0:window], trig)
        toeplitz_mat[eid] = toep_mat
        full_toep.append(toep_mat)

    # Concatenate toeplitz
    full_toep = np.concatenate(full_toep)

    # least square estimation
    predictor = np.dot(linalg.pinv(np.dot(full_toep, full_toep.T)), full_toep)
    all_evokeds = np.dot(predictor, data.T)
    all_evokeds = np.vsplit(all_evokeds, len(event_id))

    # parse evoked response
    evoked_data = dict()
    for idx, eid in enumerate(event_id):
        evoked_data[eid] = all_evokeds[idx].T

    return evoked_data, toeplitz_mat


def _check_overlapp(epochs):
    """check if events are overlapped."""
    isi = np.diff(epochs.events[:, 0])
    window = int((epochs.tmax - epochs.tmin) * epochs.info['sfreq'])
    # Events are overlapped if the minimal inter-stimulus interval is smaller
    # than the time window.
    return isi.min() < window


def _construct_signal_from_epochs(events, sfreq, tmin, tmax, epochs_data):
    """Reconstruct pseudo continuous signal from epochs."""
    start = (np.min(events[:, 0]) + int(tmin * sfreq))
    stop = (np.max(events[:, 0]) + int(tmax * sfreq) + 1)

    n_samples = stop - start
    n_epochs, n_channels, n_times = epochs_data.shape
    events_pos = events[:, 0] - events[0, 0]

    data = np.zeros((n_channels, n_samples))
    for idx in range(n_epochs):
        onset = events_pos[idx]
        offset = onset + n_times
        data[:, onset:offset] = epochs_data[idx]

    return data


def _least_square_evoked_epochs(epochs_data, events, event_id, tmin,
                                tmax, sfreq, return_toeplitz=False):
    """Least square estimation of evoked response from a Epochs instance.

    Parameters
    ----------
    epochs : Epochs instance
        An instance of Epochs.
    return_toeplitz : bool (default False)
        If true, compute the toeplitz matrix.

    Returns
    -------
    evokeds : dict of evoked instance
        An dict of evoked instance for each event type in epochs.event_id.
    toeplitz : dict of ndarray
        If return_toeplitz is true, return the toeplitz matrix for each event
        type in epochs.event_id.
    """

    events = events.copy()
    events[:, 0] -= events[0, 0] + int(tmin * sfreq)
    data = _construct_signal_from_epochs(events, sfreq, tmin, tmax,
                                         epochs_data)
    evoked_data, toeplitz = _least_square_evoked(data, events, event_id,
                                                 tmin=tmin, tmax=tmax,
                                                 sfreq=sfreq)
    evokeds = dict()
    for name, data in evoked_data.items():
        evokeds[name] = data

    if return_toeplitz:
        return evokeds, toeplitz

    return evokeds


class _Xdawn(object):
    """Private class containing common functions for Xdawn and
       XdawnTransformer.
    """
    def __init__(self, n_components, signal_cov, reg):
        self.n_components = n_components
        self.signal_cov = signal_cov
        self.reg = reg

    def _set_signal_cov(self, epochs_data):
        if self.signal_cov is None:
            sig_data = np.hstack(epochs_data)
            self.signal_cov_ = _regularized_covariance(sig_data, self.reg)
        elif isinstance(self.signal_cov, Covariance):
            self.signal_cov_ = self.signal_cov.data
        elif isinstance(self.signal_cov, np.ndarray):
            self.signal_cov_ = self.signal_cov
        else:
            raise ValueError('signal_cov must be None, a covariance instance '
                             'or a ndarray')

    def _epochs_data_dict(self, epochs_data, y, event_id):
        """Select index of same event_id from epochs_data and make a dict."""
        # TODO: Implement pythonic way of writing this.
        epochs_data_dict = dict()
        for eid in event_id:
            index = list()
            for idx in range(len(y)):
                if y[idx] == event_id[eid]:
                    index.append(idx)
            epochs_data_dict[eid] = epochs_data[index]

        return epochs_data_dict

    def _fit_xdawn(self, epochs_data, y, event_id, events=None, tmin=None,
                   tmax=None, sfreq=None, baseline=None):
        self.filters_ = dict()
        self.patterns_ = dict()
        self.evokeds_cov_ = dict()

        if self.correct_overlap:
            if baseline is not None:
                raise ValueError('Baseline correction must be None if overlap '
                                 'correction activated')
            evokeds, toeplitz =
            _least_square_evoked_epochs(epochs_data, events, event_id, tmin,
                                        tmax, sfreq, return_toeplitz=True)
        else:
            epochs_data_dict = self._epochs_data_dict(epochs_data, y,
                                                      event_id)
            evokeds = dict()
            toeplitz = dict()
            for eid in event_id:<commit_msg>pep8 issue<commit_after>"""Xdawn implementation."""
# Authors: <NAME> <<EMAIL>>
#
# License: BSD (3-clause)

import copy as cp

import numpy as np
from scipy import linalg

from ..io.base import _BaseRaw
from ..epochs import _BaseEpochs
from .. import Covariance, EvokedArray, Evoked, EpochsArray
from ..io.pick import pick_types
from .ica import _get_fast_dot
from ..utils import logger
from ..cov import _regularized_covariance
from ..channels.channels import ContainsMixin


def _least_square_evoked(data, events, event_id, tmin, tmax, sfreq):
    """Least square estimation of evoked response from data.

    Parameters
    ----------
    data : ndarray, shape (n_channels, n_times)
        The data to estimates evoked
    events : ndarray, shape (n_events, 3)
        The events typically returned by the read_events function.
        If some events don't match the events of interest as specified
        by event_id, they will be ignored.
    event_id : dict
        The id of the events to consider
    tmin : float
        Start time before event.
    tmax : float
        End time after event.
    sfreq : float
        Sampling frequency.

    Returns
    -------
    evokeds_data : dict of ndarray
        A dict of evoked data for each event type in event_id.
    toeplitz : dict of ndarray
        A dict of toeplitz matrix for each event type in event_id.
    """
    nmin = int(tmin * sfreq)
    nmax = int(tmax * sfreq)

    window = nmax - nmin
    n_samples = data.shape[1]
    toeplitz_mat = dict()
    full_toep = list()
    for eid in event_id:
        # select events by type
        ix_ev = events[:, -1] == event_id[eid]

        # build toeplitz matrix
        trig = np.zeros((n_samples, 1))
        ix_trig = (events[ix_ev, 0]) + nmin
        trig[ix_trig] = 1
        toep_mat = linalg.toeplitz(trig[0:window], trig)
        toeplitz_mat[eid] = toep_mat
        full_toep.append(toep_mat)

    # Concatenate toeplitz
    full_toep = np.concatenate(full_toep)

    # least square estimation
    predictor = np.dot(linalg.pinv(np.dot(full_toep, full_toep.T)), full_toep)
    all_evokeds = np.dot(predictor, data.T)
    all_evokeds = np.vsplit(all_evokeds, len(event_id))

    # parse evoked response
    evoked_data = dict()
    for idx, eid in enumerate(event_id):
        evoked_data[eid] = all_evokeds[idx].T

    return evoked_data, toeplitz_mat


def _check_overlapp(epochs):
    """check if events are overlapped."""
    isi = np.diff(epochs.events[:, 0])
    window = int((epochs.tmax - epochs.tmin) * epochs.info['sfreq'])
    # Events are overlapped if the minimal inter-stimulus interval is smaller
    # than the time window.
    return isi.min() < window


def _construct_signal_from_epochs(events, sfreq, tmin, tmax, epochs_data):
    """Reconstruct pseudo continuous signal from epochs."""
    start = (np.min(events[:, 0]) + int(tmin * sfreq))
    stop = (np.max(events[:, 0]) + int(tmax * sfreq) + 1)

    n_samples = stop - start
    n_epochs, n_channels, n_times = epochs_data.shape
    events_pos = events[:, 0] - events[0, 0]

    data = np.zeros((n_channels, n_samples))
    for idx in range(n_epochs):
        onset = events_pos[idx]
        offset = onset + n_times
        data[:, onset:offset] = epochs_data[idx]

    return data


def _least_square_evoked_epochs(epochs_data, events, event_id, tmin,
                                tmax, sfreq, return_toeplitz=False):
    """Least square estimation of evoked response from a Epochs instance.

    Parameters
    ----------
    epochs : Epochs instance
        An instance of Epochs.
    return_toeplitz : bool (default False)
        If true, compute the toeplitz matrix.

    Returns
    -------
    evokeds : dict of evoked instance
        An dict of evoked instance for each event type in epochs.event_id.
    toeplitz : dict of ndarray
        If return_toeplitz is true, return the toeplitz matrix for each event
        type in epochs.event_id.
    """

    events = events.copy()
    events[:, 0] -= events[0, 0] + int(tmin * sfreq)
    data = _construct_signal_from_epochs(events, sfreq, tmin, tmax,
                                         epochs_data)
    evoked_data, toeplitz = _least_square_evoked(data, events, event_id,
                                                 tmin=tmin, tmax=tmax,
                                                 sfreq=sfreq)
    evokeds = dict()
    for name, data in evoked_data.items():
        evokeds[name] = data

    if return_toeplitz:
        return evokeds, toeplitz

    return evokeds


class _Xdawn(object):
    """Private class containing common functions for Xdawn and
       XdawnTransformer.
    """
    def __init__(self, n_components, signal_cov, reg):
        self.n_components = n_components
        self.signal_cov = signal_cov
        self.reg = reg

    def _set_signal_cov(self, epochs_data):
        if self.signal_cov is None:
            sig_data = np.hstack(epochs_data)
            self.signal_cov_ = _regularized_covariance(sig_data, self.reg)
        elif isinstance(self.signal_cov, Covariance):
            self.signal_cov_ = self.signal_cov.data
        elif isinstance(self.signal_cov, np.ndarray):
            self.signal_cov_ = self.signal_cov
        else:
            raise ValueError('signal_cov must be None, a covariance instance '
                             'or a ndarray')

    def _epochs_data_dict(self, epochs_data, y, event_id):
        """Select index of same event_id from epochs_data and make a dict."""
        # TODO: Implement pythonic way of writing this.
        epochs_data_dict = dict()
        for eid in event_id:
            index = list()
            for idx in range(len(y)):
                if y[idx] == event_id[eid]:
                    index.append(idx)
            epochs_data_dict[eid] = epochs_data[index]

        return epochs_data_dict

    def _fit_xdawn(self, epochs_data, y, event_id, events=None, tmin=None,
                   tmax=None, sfreq=None, baseline=None):
        self.filters_ = dict()
        self.patterns_ = dict()
        self.evokeds_cov_ = dict()

        if self.correct_overlap:
            if baseline is not None:
                raise ValueError('Baseline correction must be None if overlap '
                                 'correction activated')
            evokeds, toeplitz = _least_square_evoked_epochs(epochs_data,
                                                            events, event_id,
                                                            tmin, tmax, sfreq,
                                                            return_toeplitz=(
                                                                True))
        else:
            epochs_data_dict = self._epochs_data_dict(epochs_data, y,
                                                      event_id)
            evokeds = dict()
            toeplitz = dict()
            for eid in event_id:
{'fraction_non_alphanumeric': 0.08824391673995896, 'fraction_numerical': 0.002711814717091762, 'mean_word_length': 2.111744583808438, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 29, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5739189', 'n_tokens_mistral': 4267, 'n_tokens_neox': 4148, 'n_words': 1179}