<commit_before>import golem
import numpy as np
import scipy
import logging
import psychic
from matplotlib.mlab import specgram

def baseline(data, baseline_period=None):
    '''
    For each channel, calculate and remove the baseline. The baseline is the
    mean signal calculated over a certain time period.

    Parameters
    ----------
    baseline_period : tuple (int, int)
        The start (inclusive) and end (exclusive) indices of the period to
        calculate the baseline over. Values are given in samples.

    Returns
    -------
    d : :class:`golem.DataSet`
        The baselined trials.
    '''
    if baseline_period:
        assert len(baseline_period) == 2, 'Specify a begin and end point for the baseline period (in samples)'
    else:
        baseline_period = (0, data.ninstances)

    assert data.ndX.ndim <= 3

    if data.ndX.ndim == 2:
        num_samples = data.ninstances
        X = data.X - np.tile( np.mean(data.X[:,baseline_period[0]:baseline_period[1]], axis=1).T, (num_samples, 1) ).T

    else:
        num_samples = data.ndX.shape[1]
        ndX = np.zeros(data.ndX.shape, dtype=data.X.dtype)

        for i in range(data.ninstances):
            ndX[:,:,i] = data.ndX[:,:,i] - np.tile( np.mean(data.ndX[:,baseline_period[0]:baseline_period[1],i], axis=1).T, (num_samples, 1) ).T
        X = ndX.reshape(data.X.shape)

    return golem.DataSet(X=X, default=data)

def erp(data, classes=None, enforce_equal_n=True):
    '''
    For each class, calculate the Event Related Potential by averaging the 
    corresponding trials. Note: no baselining is performed, see
    :func:`psychic.baseline`.

    Parameters
    ----------
    data : :class:`golem.DataSet`
        The trials
    classes: list (optional)
        When specified, the ERP is only calculated for the classes with the
        given indices.
    enforce_equal_n : bool (default=True)
        When set, each ERP is calculated by averaging the same number of
        trials. For example, if class1 has m and class2 has n trials and m > n.
        The ERP for class1 will be calculated by taking n random trials from
        class1.
    
    Returns
    -------
    d : :class:`golem.DataSet`
        A golem DataSet containing for each class the ERP. 

        - ``d.ndX``: [channels x samples x classes]
        - ``d.Y``: The class labels. Each class has one instance (one ERP).
    '''
    assert data.ndX.ndim > 2

    if classes == None or len(classes) == 0:
        # Take all classes with >0 instances
        classes = [cl for cl in range(data.nclasses) if data.ninstances_per_class[cl] > 0]
    assert len(classes) > 0, 'No valid classes specified and no classes found with >0 instances'

    num_trials = np.min( np.array(data.ninstances_per_class)[classes] )
    assert num_trials > 0, 'For one or more classes there are no instances!'

    # Calculate ERP
    erp = np.zeros(data.ndX.shape[:-1] + (len(classes),))
    for i,cl in enumerate(classes):
        trials = data.get_class(cl).ndX

        if enforce_equal_n:
            # Enforce an equal number of trials for all classes. Picking them at random.
            # Otherwise the ERPs will be skewed, simply because a different number of trials are averaged.
            idx = range(trials.shape[-1])[:num_trials]
            np.random.shuffle(idx)
            erp[...,i] = np.mean(trials[...,idx], axis=trials.ndim-1)
        else:
            erp[...,i] = np.mean(trials, axis=trials.ndim-1)

    X = erp.reshape(-1, len(classes))
    Y = golem.helpers.to_one_of_n(classes).astype(np.bool)
    I = np.atleast_2d(classes)
    feat_shape = data.ndX.shape[:-1]
    cl_lab = [lab for i,lab in enumerate(data.cl_lab) if i in classes]

    return golem.DataSet(X=X, Y=Y, I=I, feat_shape=feat_shape, cl_lab=cl_lab, default=data)

def ttest(data, classes=[0, 1], shuffle=True):
    '''
    Calculate ttests between two classes for each channel and each sample. If
    one class has more trials than the other, random trials will be taken to
    ensure an equal number of trials.

    Parameters
    ----------

    classes : list (default=[0, 1])
        The indices of the classes to compare.
    shuffle : bool (default=False)
        When set, trials will be shuffled prior to comparison.

    Returns
    -------
    t-values: [channels x samples] array
        The t-values 

    p-values: [channels x samples] array
        The p-values 
    '''
    assert data.nd_xs.ndim == 3
    assert data.nclasses >= 2, ('Data must contain at least two classes ',
                                'otherwise there is nothing to compare.')

    num_trials = np.min( np.array(data.ninstances_per_class)[classes] )

    c1 = data.ndX[..., data.Y[classes[0],:].astype(np.bool)]
    if shuffle:
        np.random.shuffle(c1)
    c1 = c1[..., :num_trials]

    c2 = data.ndX[..., data.Y[classes[1],:].astype(np.bool)]
    if shuffle:
        np.random.shuffle(c2)
    c2 = c2[..., :num_trials]

    return scipy.stats.ttest_ind(c1, c2, axis=2)

def random_groups(d, size):
    '''
    For each class, form groups of random trials of the given size.

    Parameters
    ----------
    d : :class:`golem.DataSet`
        The trials.
    size : int
        Size of the groups to make.

    Returns
    -------
    d : :class:`golem.DataSet`
        The grouped data. ``d.ndX`` is [channels x samples x trials x groups]
    '''

    d_trials = None
    idxs = []

    for cl in range(d.nclasses):
        idx_cl = np.flatnonzero(d.Y[cl,:])
        ninstances = len(idx_cl)
        ngroups = int(ninstances / size)

        idx = idx_cl[np.random.permutation(ninstances)[:ngroups*size]].reshape(size,-1)
        d_grouped = golem.DataSet(
            X = d.ndX[:,:,idx].reshape(-1, ngroups),
            Y = d.Y[:,idx[0,:]],
            I = d.I[:,idx[0,:]],
            feat_shape = d.feat_shape + (size,),
            feat_dim_lab = d.feat_dim_lab + ['trials'] if d.feat_dim_lab else None,
            feat_nd_lab = d.feat_nd_lab + [range(size)] if d.feat_nd_lab else None,
            default = d
        )

        if d_trials == None:
            d_trials = d_grouped
        else:
            d_trials += d_grouped

        idxs.append(idx)

    return (d_trials, np.hstack(idxs))

def reject_trials(d, cutoff=100, time_range=None):
    '''
    Reject trials by thresholding the absolute amplitude. 

    Parameters
    ----------

    d : :class:`golem.DataSet`
        The dataset to filter.
    cutoff : float (default=100)
        Any trials with a feature larger than this value are rejected.
        Alternatively, a list containing cutoff values for each channel can be
        specified. 
    time_range : tuple (default=all samples)
        Range (begin, end) for which to apply the thresholding. Values are
        given as sample indices.

    Returns
    -------
    d : :class:`golem.DataSet`
        Filtered dataset.<commit_msg>Added ability to `random_groups` to do blowup operations. Also increased efficiency of the algorithm.<commit_after>import golem
import numpy as np
import scipy
import logging
import psychic
from matplotlib.mlab import specgram

def baseline(data, baseline_period=None):
    '''
    For each channel, calculate and remove the baseline. The baseline is the
    mean signal calculated over a certain time period.

    Parameters
    ----------
    baseline_period : tuple (int, int)
        The start (inclusive) and end (exclusive) indices of the period to
        calculate the baseline over. Values are given in samples.

    Returns
    -------
    d : :class:`golem.DataSet`
        The baselined trials.
    '''
    if baseline_period:
        assert len(baseline_period) == 2, 'Specify a begin and end point for the baseline period (in samples)'
    else:
        baseline_period = (0, data.ninstances)

    assert data.ndX.ndim <= 3

    if data.ndX.ndim == 2:
        num_samples = data.ninstances
        X = data.X - np.tile( np.mean(data.X[:,baseline_period[0]:baseline_period[1]], axis=1).T, (num_samples, 1) ).T

    else:
        num_samples = data.ndX.shape[1]
        ndX = np.zeros(data.ndX.shape, dtype=data.X.dtype)

        for i in range(data.ninstances):
            ndX[:,:,i] = data.ndX[:,:,i] - np.tile( np.mean(data.ndX[:,baseline_period[0]:baseline_period[1],i], axis=1).T, (num_samples, 1) ).T
        X = ndX.reshape(data.X.shape)

    return golem.DataSet(X=X, default=data)

def erp(data, classes=None, enforce_equal_n=True):
    '''
    For each class, calculate the Event Related Potential by averaging the 
    corresponding trials. Note: no baselining is performed, see
    :func:`psychic.baseline`.

    Parameters
    ----------
    data : :class:`golem.DataSet`
        The trials
    classes: list (optional)
        When specified, the ERP is only calculated for the classes with the
        given indices.
    enforce_equal_n : bool (default=True)
        When set, each ERP is calculated by averaging the same number of
        trials. For example, if class1 has m and class2 has n trials and m > n.
        The ERP for class1 will be calculated by taking n random trials from
        class1.
    
    Returns
    -------
    d : :class:`golem.DataSet`
        A golem DataSet containing for each class the ERP. 

        - ``d.ndX``: [channels x samples x classes]
        - ``d.Y``: The class labels. Each class has one instance (one ERP).
    '''
    assert data.ndX.ndim > 2

    if classes == None or len(classes) == 0:
        # Take all classes with >0 instances
        classes = [cl for cl in range(data.nclasses) if data.ninstances_per_class[cl] > 0]
    assert len(classes) > 0, 'No valid classes specified and no classes found with >0 instances'

    num_trials = np.min( np.array(data.ninstances_per_class)[classes] )
    assert num_trials > 0, 'For one or more classes there are no instances!'

    # Calculate ERP
    erp = np.zeros(data.ndX.shape[:-1] + (len(classes),))
    for i,cl in enumerate(classes):
        trials = data.get_class(cl).ndX

        if enforce_equal_n:
            # Enforce an equal number of trials for all classes. Picking them at random.
            # Otherwise the ERPs will be skewed, simply because a different number of trials are averaged.
            idx = range(trials.shape[-1])[:num_trials]
            np.random.shuffle(idx)
            erp[...,i] = np.mean(trials[...,idx], axis=trials.ndim-1)
        else:
            erp[...,i] = np.mean(trials, axis=trials.ndim-1)

    X = erp.reshape(-1, len(classes))
    Y = golem.helpers.to_one_of_n(classes).astype(np.bool)
    I = np.atleast_2d(classes)
    feat_shape = data.ndX.shape[:-1]
    cl_lab = [lab for i,lab in enumerate(data.cl_lab) if i in classes]

    return golem.DataSet(X=X, Y=Y, I=I, feat_shape=feat_shape, cl_lab=cl_lab, default=data)

def ttest(data, classes=[0, 1], shuffle=True):
    '''
    Calculate ttests between two classes for each channel and each sample. If
    one class has more trials than the other, random trials will be taken to
    ensure an equal number of trials.

    Parameters
    ----------

    classes : list (default=[0, 1])
        The indices of the classes to compare.
    shuffle : bool (default=False)
        When set, trials will be shuffled prior to comparison.

    Returns
    -------
    t-values: [channels x samples] array
        The t-values 

    p-values: [channels x samples] array
        The p-values 
    '''
    assert data.nd_xs.ndim == 3
    assert data.nclasses >= 2, ('Data must contain at least two classes ',
                                'otherwise there is nothing to compare.')

    num_trials = np.min( np.array(data.ninstances_per_class)[classes] )

    c1 = data.ndX[..., data.Y[classes[0],:].astype(np.bool)]
    if shuffle:
        np.random.shuffle(c1)
    c1 = c1[..., :num_trials]

    c2 = data.ndX[..., data.Y[classes[1],:].astype(np.bool)]
    if shuffle:
        np.random.shuffle(c2)
    c2 = c2[..., :num_trials]

    return scipy.stats.ttest_ind(c1, c2, axis=2)

def random_groups(d, group_size, groups_per_class=None, mean=False):
    '''
    For each class, form groups of random trials of the given size.

    Parameters
    ----------
    d : :class:`golem.DataSet`
        The trials.
    group_size : int
        Size of the groups to make.
    groups_per_class : int (default=ninstances_per_class / group_size)
        Number of groups to make per class.
    mean : bool (default=False)
        Return group means (i.e. ERPs) instead of groups

    Returns
    -------
    d : :class:`golem.DataSet`
        The grouped data. ``d.ndX`` is [channels x samples x trials x groups] or
        if mean==True, ``d.ndX`` is [channels x samples x groups]
    '''

    d_trials = []
    idxs = []

    if groups_per_class == None:
        groups_per_class = np.min(d.ninstances_per_class) / group_size

    for cl in range(d.nclasses):
        idx_cl = np.flatnonzero(d.Y[cl,:])
        ninstances = len(idx_cl)

        groups_to_go = groups_per_class
        while groups_to_go > 0:
            ngroups = min(int(ninstances / group_size), groups_to_go)

            idx = idx_cl[np.random.permutation(ninstances)[:ngroups*group_size]].reshape(group_size,-1)

            ndX = d.ndX[:,:,idx]
            if mean:
                ndX = np.mean(d.ndX[:,:,idx], axis=2)

            if mean:
                feat_dim_lab = d.feat_dim_lab
            else:
                feat_dim_lab = d.feat_dim_lab + ['trials'] if d.feat_dim_lab else None

            if mean:
                feat_nd_lab = d.feat_nd_lab
            else:
                feat_nd_lab = d.feat_nd_lab + [range(group_size)] if d.feat_nd_lab else None

            Y = d.Y[:,idx[0,:]]
            I = d.I[:,idx[0,:]]

            d_grouped = golem.DataSet(
                ndX=ndX, Y=Y, I=I, feat_dim_lab=feat_dim_lab,
                feat_nd_lab=feat_nd_lab, default = d
            )

            d_trials.append(d_grouped)
            idxs.append(idx)

            groups_to_go -= ngroups

    if len(d_trials) == 0:
        return ([], [])
    else:
        return (golem.concatenate(d_trials, ignore_index=True), np.hstack(idxs))

def reject_trials(d, cutoff=100, time_range=None):
    '''
    Reject trials by thresholding the absolute amplitude. 

    Parameters
    ----------

    d : :class:`golem.DataSet`
        The dataset to filter.
    cutoff : float (default=100)
        Any trials with a feature larger than this value are rejected.
        Alternatively, a list containing cutoff values for each channel can be
        specified. 
    time_range : tuple (default=all samples)
        Range (begin, end) for which to apply the thresholding. Values are
        given as sample indices.

    Returns
    -------
    d : :class:`golem.DataSet`
        Filtered dataset.
{'fraction_non_alphanumeric': 0.11775055830006091, 'fraction_numerical': 0.008256073627935305, 'mean_word_length': 2.3075201432408234, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 20, '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': '1308573', 'n_tokens_mistral': 4828, 'n_tokens_neox': 4678, 'n_words': 1406}