<commit_before>import numpy as np
from psychic.nodes.spatialfilter import BaseSpatialFilter, sfilter_trial
import psychic


def _calc_beamformer_snr(X1, X2, nc=1, a=1.0):
        nsamples, nchannels, ninstances1 = X1.shape
        ninstances2 = X2.shape[2]

        R1 = np.zeros((nchannels, nchannels, ninstances1))
        for i in range(ninstances1):
            cov = X1[...,i].T.dot(X1[...,i])
            R1[...,i] = cov / np.trace(cov)
        R1 = np.mean(R1, axis=2)

        R2 = np.zeros((nchannels, nchannels, ninstances2))
        for i in range(ninstances2):
            cov = X2[...,i].T.dot(X2[...,i])
            R2[...,i] = cov / np.trace(cov)
        R2 = np.mean(R2, axis=2)

        V, W = np.linalg.eig(np.linalg.pinv(a*R2).dot(R1))
        order = np.argsort(V)[::-1]
        V = V[order][:nc]
        W = np.real(W[:,order][:,:nc])

        return (V,W)

def _calc_beamformer_fc(Xs, nc=1, theta=1.0):
        nsamples, nchannels = Xs[0].shape[:2]
        ninstances = [X.shape[2] for X in Xs]
        p = [n / float(np.sum(ninstances)) for n in ninstances]

        means = [np.mean(X, axis=2) for X in Xs]
        grand_avg = np.mean(np.concatenate([m[..., np.newaxis] for m in means], axis=2), axis=2)

        S_b = np.zeros((nchannels, nchannels))
        for i in range(len(Xs)):
            diff = means[i] - grand_avg
            S_b += p[i] * diff.T.dot(diff)

        S_w = np.zeros((nchannels, nchannels))
        for i,X in enumerate(Xs):
            for k in range(ninstances[i]):
                diff = X[...,k] - means[i]
                S_w += diff.T.dot(diff)

        I = np.identity(nchannels)

        V, W = np.linalg.eig(np.linalg.pinv((I-theta).dot(S_w) + theta*I).dot(S_b))

        order = np.argsort(V)[::-1]
        V = V[order][:nc]
        W = np.real(W[:,order][:,:nc])

        return (V,W)


class BeamformerSNR(BaseSpatialFilter):
    def __init__(self, nc=1, a=1.0):
        assert 0 < a <= 1, 'Regularization parameter should be in range (0; 1]'
        BaseSpatialFilter.__init__(self, 1)
        self.nc = nc
        self.a = a

    def train_(self, d):
        assert len(d.feat_shape) == 2, 'Expecting sliced data'
        assert d.nclasses == 2, 'Expecting exactly two classes'

        X1 = d.get_class(0).ndX
        X2 = d.get_class(1).ndX
        self.V, self.W = _calc_beamformer_snr(X1, X2, self.nc, self.a)

class BeamformerFC(BaseSpatialFilter):
    def __init__(self, nc=1, theta=1.0):
        assert 0 < theta <= 1, 'Regularization parameter should be in range (0; 1]'
        BaseSpatialFilter.__init__(self, 1)
        self.nc = nc
        self.theta = theta

    def train_(self, d):
        assert len(d.feat_shape) == 2, 'Expecting sliced data'
        assert d.nclasses > 1, 'Expecting more than one class'

        Xs = [d.get_class(i).ndX for i in range(d.nclasses)]
        self.V, self.W = _calc_beamformer_fc(Xs, self.nc, self.theta)

class BeamformerCFMS(BaseSpatialFilter):
    def __init__(self, nc=2, theta=1.0):
        assert nc % 2 == 0, 'Number of components should be even'
        assert 0 < theta <= 1, 'Regularization parameter should be in range (0; 1]'
        BaseSpatialFilter.__init__(self, 1)
        self.nc = nc
        self.theta = theta

    def train_(self, d):
        assert len(d.feat_shape) == 2, 'Expecting sliced data'
        assert d.nclasses == 2, 'Expecting exactly two classes'
        nchannels = d.feat_shape[1]

        # Calculate FC beamformer
        X1 = d.get_class(0).ndX
        X2 = d.get_class(1).ndX
        W_fc,_ = _calc_beamformer_fc([X1, X2], nchannels, self.theta)

        # Apply FC beamformer
        d2 = sfilter_trial(d, W_fc)

        # Calculate SNR beamformer on the result (skip first nc/2 components)
        X1 = d2.get_class(0).ndX[:,self.nc/2:,:]
        X2 = d2.get_class(1).ndX[:,self.nc/2:,:]
        W_snr,_ = _calc_beamformer_snr(X1, X2, self.nc/2, self.theta)

        print W_fc.shape
        print W_snr.shape
        self.W = np.dot(W_fc, W_snr)<commit_msg>Fixed max-SNR and CFMS beamformers<commit_after>'''
Implements beamformer spatial filters to separate ERPs as described in [1].

[1] <NAME>, <NAME>, and <NAME>. Statistical spatial
filtering for a P300-based BCI: Test in able-bodied, and patient with cerebral
palsy and amyotrophic lateral sclerosis. Journal of Neuroscience Methods, 
195:270-281, 2011.
'''

import numpy as np
from psychic.nodes.spatialfilter import BaseSpatialFilter, sfilter_trial

def _calc_beamformer_snr(X1, X2, nc=1, a=1.0):
        nchannels, nsamples, ninstances1 = X1.shape
        ninstances2 = X2.shape[2]

        R1 = np.zeros((nchannels, nchannels, ninstances1))
        for i in range(ninstances1):
            cov = X1[...,i].dot(X1[...,i].T)
            R1[...,i] = cov / np.trace(cov)
        R1 = np.mean(R1, axis=2)

        R2 = np.zeros((nchannels, nchannels, ninstances2))
        for i in range(ninstances2):
            cov = X2[...,i].dot(X2[...,i].T)
            R2[...,i] = cov / np.trace(cov)
        R2 = np.mean(R2, axis=2)

        V, W = np.linalg.eig(np.linalg.pinv(a*R2).dot(R1))
        order = np.argsort(V)[::-1]
        V = V[order][:nc]
        W = np.real(W[:,order][:,:nc])

        return (V,W)

def _calc_beamformer_fc(Xs, nc=1, theta=1.0):
        nchannels, nsamples = Xs[0].shape[:2]
        ninstances = [X.shape[2] for X in Xs]
        p = [n / float(np.sum(ninstances)) for n in ninstances]

        means = [np.mean(X, axis=2) for X in Xs]
        grand_avg = np.mean(np.concatenate([m[..., np.newaxis] for m in means], axis=2), axis=2)

        S_b = np.zeros((nchannels, nchannels))
        for i in range(len(Xs)):
            diff = means[i] - grand_avg
            S_b += p[i] * diff.dot(diff.T)

        S_w = np.zeros((nchannels, nchannels))
        for i,X in enumerate(Xs):
            for k in range(ninstances[i]):
                diff = X[...,k] - means[i]
                S_w += diff.dot(diff.T)

        I = np.identity(nchannels)

        V, W = np.linalg.eig(np.linalg.pinv((I-theta).dot(S_w) + theta*I).dot(S_b))

        order = np.argsort(V)[::-1]
        V = V[order][:nc]
        W = np.real(W[:,order][:,:nc])

        return (V,W)


class BeamformerSNR(BaseSpatialFilter):
    '''
    A spatial ERP filter that aims to separate two classes. It optimizes the
    signal to noise ratio where class 0 is taken as signal and class 1 as noise.
    '''
    def __init__(self, nc=1, theta=1.0):
        '''
        Creates an max-SNR beamformer that will generate nc components. A
        regularization parameter theta can be supplied (0..1) to prevent
        overfitting.
        '''
        assert 0 < theta <= 1, 'Regularization parameter should be in range (0; 1]'
        BaseSpatialFilter.__init__(self, 1)
        self.nc = nc
        self.theta = theta

    def train_(self, d):
        assert len(d.feat_shape) == 2, 'Expecting sliced data'
        assert d.nclasses == 2, 'Expecting exactly two classes'

        X1 = d.get_class(0).ndX
        X2 = d.get_class(1).ndX
        self.V, self.W = _calc_beamformer_snr(X1, X2, self.nc, self.theta)

class BeamformerFC(BaseSpatialFilter):
    '''
    A spatial ERP filter that aims to separate two classes. It optimizes the
    Fisher's criterion, so it increases the separation between classes while
    minimizing the variance within a class.
    '''
    def __init__(self, nc=1, theta=1.0):
        '''
        Creates an FC beamformer that will generate nc components. A
        regularization parameter theta can be supplied (0..1) to prevent
        overfitting.
        '''
        assert 0 < theta <= 1, 'Regularization parameter should be in range (0; 1]'
        BaseSpatialFilter.__init__(self, 1)
        self.nc = nc
        self.theta = theta

    def train_(self, d):
        assert len(d.feat_shape) == 2, 'Expecting sliced data'
        assert d.nclasses > 1, 'Expecting more than one class'

        Xs = [d.get_class(i).ndX for i in range(d.nclasses)]
        self.V, self.W = _calc_beamformer_fc(Xs, self.nc, self.theta)

class BeamformerCFMS(BaseSpatialFilter):
    '''
    A combination of the BeamformerFC and BeamformerSNR classes that combines
    both methods in a suboptimum way. First BeamformerFC is run to generate the
    first nc/2 components. Then, BeamformerSNR is run on the FC result to
    generate the last nc/2 components.
    '''
    def __init__(self, nc=2, theta=1.0):
        '''
        Creates an CFMS beamformer that will generate nc components, where half
        of the components are supplied by an FC beamformer and half are supplied
        by an max-SNR beamformer. NC should therefore alsways be an even number.
        A regularization parameter theta can be supplied (0..1) to prevent
        overfitting. 
        '''
        assert nc % 2 == 0, 'Number of components should be even'
        assert 0 < theta <= 1, 'Regularization parameter should be in range (0; 1]'
        BaseSpatialFilter.__init__(self, 1)
        self.nc = nc
        self.theta = theta

    def train_(self, d):
        assert len(d.feat_shape) == 2, 'Expecting sliced data'
        assert d.nclasses == 2, 'Expecting exactly two classes'
        nchannels = d.feat_shape[0]

        # Calculate FC beamformer
        X1 = d.get_class(0).ndX
        X2 = d.get_class(1).ndX
        _, W_fc = _calc_beamformer_fc([X1, X2], nchannels, self.theta)

        # Apply FC beamformer
        d2 = sfilter_trial(d, W_fc)

        # Calculate SNR beamformer on the result (skip first nc/2 components)
        X1 = d2.get_class(0).ndX[:,self.nc/2:,:]
        X2 = d2.get_class(1).ndX[:,self.nc/2:,:]
        _, W_snr = _calc_beamformer_snr(X1, X2, self.nc/2, self.theta)

        I = np.identity(nchannels)

        # Construct filter that will take the first nc/2 FC components applied
        # to the EEG data, and then the first nc/2 SNR components applied to the
        # FC filtered data.
        self.W = np.dot(W_fc, np.concatenate((I[:,:self.nc/2], W_snr), axis=1))
{'fraction_non_alphanumeric': 0.1312868713128687, 'fraction_numerical': 0.022997700229977002, 'mean_word_length': 2.318513603185136, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2822814', 'n_tokens_mistral': 3697, 'n_tokens_neox': 3505, 'n_words': 940}