<s> """<STR_LIT>""" <EOL> import copy as cp <EOL> import numpy as np <EOL> from scipy import linalg <EOL> from . . io . base import _BaseRaw <EOL> from . . epochs import _BaseEpochs <EOL> from . . import Covariance , EvokedArray , Evoked , EpochsArray <EOL> from . . io . pick import pick_types <EOL> from . ica import _get_fast_dot <EOL> from . . utils import logger <EOL> from . . decoding . mixin import TransformerMixin <EOL> from . . cov import _regularized_covariance <EOL> from . . channels . channels import ContainsMixin <EOL> def _least_square_evoked ( data , events , event_id , tmin , tmax , sfreq ) : <EOL> """<STR_LIT>""" <EOL> nmin = int ( tmin * sfreq ) <EOL> nmax = int ( tmax * sfreq ) <EOL> window = nmax - nmin <EOL> n_samples = data . shape [ <NUM_LIT:1> ] <EOL> toeplitz_mat = dict ( ) <EOL> full_toep = list ( ) <EOL> for eid in event_id : <EOL> ix_ev = events [ : , - <NUM_LIT:1> ] == event_id [ eid ] <EOL> trig = np . zeros ( ( n_samples , <NUM_LIT:1> ) ) <EOL> ix_trig = ( events [ ix_ev , <NUM_LIT:0> ] ) + nmin <EOL> trig [ ix_trig ] = <NUM_LIT:1> <EOL> toep_mat = linalg . toeplitz ( trig [ <NUM_LIT:0> : window ] , trig ) <EOL> toeplitz_mat [ eid ] = toep_mat <EOL> full_toep . append ( toep_mat ) <EOL> full_toep = np . concatenate ( full_toep ) <EOL> predictor = np . dot ( linalg . pinv ( np . dot ( full_toep , full_toep . T ) ) , full_toep ) <EOL> all_evokeds = np . dot ( predictor , data . T ) <EOL> all_evokeds = np . vsplit ( all_evokeds , len ( event_id ) ) <EOL> evoked_data = dict ( ) <EOL> for idx , eid in enumerate ( event_id ) : <EOL> evoked_data [ eid ] = all_evokeds [ idx ] . T <EOL> return evoked_data , toeplitz_mat <EOL> def _check_overlapp ( epochs ) : <EOL> """<STR_LIT>""" <EOL> isi = np . diff ( epochs . events [ : , <NUM_LIT:0> ] ) <EOL> window = int ( ( epochs . tmax - epochs . tmin ) * epochs . info [ '<STR_LIT>' ] ) <EOL> return isi . min ( ) < window <EOL> def _construct_signal_from_epochs ( epochs ) : <EOL> """<STR_LIT>""" <EOL> start = ( np . min ( epochs . events [ : , <NUM_LIT:0> ] ) + <EOL> int ( epochs . tmin * epochs . info [ '<STR_LIT>' ] ) ) <EOL> stop = ( np . max ( epochs . events [ : , <NUM_LIT:0> ] ) + <EOL> int ( epochs . tmax * epochs . info [ '<STR_LIT>' ] ) + <NUM_LIT:1> ) <EOL> n_samples = stop - start <EOL> epochs_data = epochs . get_data ( ) <EOL> n_epochs , n_channels , n_times = epochs_data . shape <EOL> events_pos = epochs . events [ : , <NUM_LIT:0> ] - epochs . events [ <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> data = np . zeros ( ( n_channels , n_samples ) ) <EOL> for idx in range ( n_epochs ) : <EOL> onset = events_pos [ idx ] <EOL> offset = onset + n_times <EOL> data [ : , onset : offset ] = epochs_data [ idx ] <EOL> return data <EOL> def least_square_evoked ( epochs , return_toeplitz = False ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( epochs , _BaseEpochs ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> events = epochs . events . copy ( ) <EOL> events [ : , <NUM_LIT:0> ] -= events [ <NUM_LIT:0> , <NUM_LIT:0> ] + int ( epochs . tmin * epochs . info [ '<STR_LIT>' ] ) <EOL> data = _construct_signal_from_epochs ( epochs ) <EOL> evoked_data , toeplitz = _least_square_evoked ( data , events , epochs . event_id , <EOL> tmin = epochs . tmin , <EOL> tmax = epochs . tmax , <EOL> sfreq = epochs . info [ '<STR_LIT>' ] ) <EOL> evokeds = dict ( ) <EOL> info = cp . deepcopy ( epochs . info ) <EOL> for name , data in evoked_data . items ( ) : <EOL> n_events = len ( events [ events [ : , <NUM_LIT:2> ] == epochs . event_id [ name ] ] ) <EOL> evoked = EvokedArray ( data , info , tmin = epochs . tmin , <EOL> comment = name , nave = n_events ) <EOL> evokeds [ name ] = evoked <EOL> if return_toeplitz : <EOL> return evokeds , toeplitz <EOL> return evokeds <EOL> class Xdawn ( TransformerMixin , ContainsMixin ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , n_components = <NUM_LIT:2> , signal_cov = None , correct_overlap = '<STR_LIT>' , <EOL> reg = None ) : <EOL> """<STR_LIT>""" <EOL> self . n_components = n_components <EOL> self . signal_cov = signal_cov <EOL> self . reg = reg <EOL> self . filters_ = dict ( ) <EOL> self . patterns_ = dict ( ) <EOL> self . evokeds_ = dict ( ) <EOL> if correct_overlap not in [ '<STR_LIT>' , True , False ] : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> self . correct_overlap = correct_overlap <EOL> def fit ( self , epochs , y = None ) : <EOL> """<STR_LIT>""" <EOL> if self . correct_overlap == '<STR_LIT>' : <EOL> self . correct_overlap = _check_overlapp ( epochs ) <EOL> if self . signal_cov is None : <EOL> if self . correct_overlap : <EOL> sig_data = _construct_signal_from_epochs ( epochs ) <EOL> else : <EOL> sig_data = np . hstack ( epochs . get_data ( ) ) <EOL> self . signal_cov_ = _regularized_covariance ( sig_data , self . reg ) <EOL> elif isinstance ( self . signal_cov , Covariance ) : <EOL> self . signal_cov_ = self . signal_cov . data <EOL> elif isinstance ( self . signal_cov , np . ndarray ) : <EOL> self . signal_cov_ = self . signal_cov <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> self . evokeds_cov_ = dict ( ) <EOL> if self . correct_overlap : <EOL> if epochs . baseline is not None : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> evokeds , toeplitz = least_square_evoked ( epochs , <EOL> return_toeplitz = True ) <EOL> else : <EOL> evokeds = dict ( ) <EOL> toeplitz = dict ( ) <EOL> for eid in epochs . event_id : <EOL> evokeds [ eid ] = epochs [ eid ] . average ( ) <EOL> toeplitz [ eid ] = <NUM_LIT:1.0> <EOL> self . evokeds_ = evokeds <EOL> for eid in epochs . event_id : <EOL> data = np . dot ( evokeds [ eid ] . data , toeplitz [ eid ] ) <EOL> self . evokeds_cov_ [ eid ] = _regularized_covariance ( data , self . reg ) <EOL> for eid in epochs . event_id : <EOL> if self . signal_cov_ . shape != self . evokeds_cov_ [ eid ] . shape : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> evals , evecs = linalg . eigh ( self . evokeds_cov_ [ eid ] , <EOL> self . signal_cov_ ) <EOL> evecs = evecs [ : , np . argsort ( evals ) [ : : - <NUM_LIT:1> ] ] <EOL> evecs /= np . sqrt ( np . sum ( evecs ** <NUM_LIT:2> , axis = <NUM_LIT:0> ) ) <EOL> self . filters_ [ eid ] = evecs <EOL> self . patterns_ [ eid ] = linalg . inv ( evecs . T ) <EOL> self . ch_names = epochs . ch_names <EOL> self . exclude = list ( range ( self . n_components , len ( self . ch_names ) ) ) <EOL> self . event_id = epochs . event_id <EOL> return self <EOL> def transform ( self , epochs ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( epochs , _BaseEpochs ) : <EOL> data = epochs . get_data ( ) <EOL> elif isinstance ( epochs , np . ndarray ) : <EOL> data = epochs <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> full_filters = list ( ) <EOL> for filt in self . filters_ . values ( ) : <EOL> full_filters . append ( filt [ : , <NUM_LIT:0> : self . n_components ] ) <EOL> full_filters = np . concatenate ( full_filters , axis = <NUM_LIT:1> ) <EOL> X = np . dot ( full_filters . T , data ) <EOL> X = X . transpose ( ( <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:2> ) ) <EOL> return X <EOL> def apply ( self , inst , event_id = None , include = None , exclude = None ) : <EOL> """<STR_LIT>""" <EOL> if event_id is None : <EOL> event_id = self . event_id <EOL> if isinstance ( inst , _BaseRaw ) : <EOL> out = self . _apply_raw ( raw = inst , include = include , exclude = exclude , <EOL> event_id = event_id ) <EOL> elif isinstance ( inst , _BaseEpochs ) : <EOL> out = self . _apply_epochs ( epochs = inst , include = include , <EOL> exclude = exclude , event_id = event_id ) <EOL> elif isinstance ( inst , Evoked ) : <EOL> out = self . _apply_evoked ( evoked = inst , include = include , <EOL> exclude = exclude , event_id = event_id ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> return out <EOL> def _apply_raw ( self , raw , include , exclude , event_id ) : <EOL> """<STR_LIT>""" <EOL> if not raw . preload : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> picks = pick_types ( raw . info , meg = False , include = self . ch_names , <EOL> exclude = '<STR_LIT>' ) <EOL> raws = dict ( ) <EOL> for eid in event_id : <EOL> data = raw [ picks , : ] [ <NUM_LIT:0> ] <EOL> data = self . _pick_sources ( data , include , exclude , eid ) <EOL> raw_r = raw . copy ( ) <EOL> raw_r [ picks , : ] = data <EOL> raws [ eid ] = raw_r <EOL> return raws <EOL> def _apply_epochs ( self , epochs , include , exclude , event_id ) : <EOL> """<STR_LIT>""" <EOL> if not epochs . preload : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> picks = pick_types ( epochs . info , meg = False , ref_meg = False , <EOL> include = self . ch_names , exclude = '<STR_LIT>' ) <EOL> if len ( picks ) != len ( self . ch_names ) : <EOL> raise RuntimeError ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' % ( len ( self . ch_names ) , <EOL> len ( picks ) ) ) <EOL> epochs_dict = dict ( ) <EOL> data = np . hstack ( epochs . get_data ( ) [ : , picks ] ) <EOL> for eid in event_id : <EOL> data_r = self . _pick_sources ( data , include , exclude , eid ) <EOL> data_r = np . array ( np . split ( data_r , len ( epochs . events ) , <NUM_LIT:1> ) ) <EOL> info_r = cp . deepcopy ( epochs . info ) <EOL> epochs_r = EpochsArray ( data = data_r , info = info_r , <EOL> events = epochs . events , tmin = epochs . tmin , <EOL> event_id = epochs . event_id , verbose = False ) <EOL> epochs_r . preload = True <EOL> epochs_dict [ eid ] = epochs_r <EOL> return epochs_dict <EOL> def _apply_evoked ( self , evoked , include , exclude , event_id ) : <EOL> """<STR_LIT>""" <EOL> picks = pick_types ( evoked . info , meg = False , ref_meg = False , <EOL> include = self . ch_names , <EOL> exclude = '<STR_LIT>' ) <EOL> if len ( picks ) != len ( self . ch_names ) : <EOL> raise RuntimeError ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' % ( len ( self . ch_names ) , <EOL> len ( picks ) ) ) <EOL> data = evoked . data [ picks ] <EOL> evokeds = dict ( ) <EOL> for eid in event_id : <EOL> data_r = self . _pick_sources ( data , include , exclude , eid ) <EOL> evokeds [ eid ] = evoked . copy ( ) <EOL> evokeds [ eid ] . data [ picks ] = data_r <EOL> return evokeds <EOL> def _pick_sources ( self , data , include , exclude , eid ) : <EOL> """<STR_LIT>""" <EOL> fast_dot = _get_fast_dot ( ) <EOL> if exclude is None : <EOL> exclude = self . exclude <EOL> else : <EOL> exclude = list ( set ( list ( self . exclude ) + list ( exclude ) ) ) <EOL> logger . info ( '<STR_LIT>' ) <EOL> sources = fast_dot ( self . filters_ [ eid ] . T , data ) <EOL> if include not in ( None , [ ] ) : <EOL> mask = np . ones ( len ( sources ) , dtype = np . bool ) <EOL> mask [ np . unique ( include ) ] = False <EOL> sources [ mask ] = <NUM_LIT:0.> <EOL> logger . info ( '<STR_LIT>' % mask . sum ( ) ) <EOL> elif exclude not in ( None , [ ] ) : <EOL> exclude_ = np . unique ( exclude ) <EOL> sources [ exclude_ ] = <NUM_LIT:0.> <EOL> logger . info ( '<STR_LIT>' % len ( exclude_ ) ) <EOL> logger . info ( '<STR_LIT>' ) <EOL> data = fast_dot ( self . patterns_ [ eid ] , sources ) <EOL> return data </s>
