kraken/backend/bci_core/model.py

import numpy as np

from sklearn.linear_model import LogisticRegression
from pyriemann.tangentspace import TangentSpace
from pyriemann.preprocessing import Whitening
from sklearn.pipeline import make_pipeline
from sklearn.base import BaseEstimator, TransformerMixin
from mne.decoding import Vectorizer


class FeatureSelector(BaseEstimator, TransformerMixin):
    def __init__(self, feature_type, lfb_dim, hgs_dim):
        self.feature_type = feature_type
        self.lfb_dim = lfb_dim
        self.hgs_dim = hgs_dim

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.feature_type == 'lfb':
            return X[:, 0:self.lfb_dim, 0:self.lfb_dim].copy()
        else:
            return X[:, self.lfb_dim:self.lfb_dim + self.hgs_dim, self.lfb_dim:self.lfb_dim + self.hgs_dim].copy()


class ChannelScaler(BaseEstimator, TransformerMixin):
    def __init__(self, norm_axis=(0, 2)):
        self.channel_mean_ = None
        self.channel_std_ = None
        self.norm_axis=norm_axis

    def fit(self, X, y=None):
        '''

        :param X: 3d array with shape (n_epochs, n_channels, n_times)
        :param y:
        :return:
        '''
        self.channel_mean_ = np.mean(X, axis=self.norm_axis, keepdims=True)
        self.channel_std_ = np.std(X, axis=self.norm_axis, keepdims=True)
        return self

    def transform(self, X, y=None):
        X = X.copy()
        X -= self.channel_mean_
        X /= self.channel_std_
        return X


def riemann_model(C=1.):
    return make_pipeline(
        Whitening(metric='riemann', dim_red={'expl_var': 0.99}),
        TangentSpace(),
        LogisticRegression(C=C)
    )


def baseline_model(C=1.):
    return make_pipeline(
        ChannelScaler(),
        Vectorizer(),
        LogisticRegression(C=C)
    )