braindecode 0.5.1__py3-none-any.whl

braindecode/util.py

@@ -0,0 +1,362 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+
+
+import os
+import glob
+import random
+
+import numpy as np
+import mne
+import h5py
+import torch
+from sklearn.utils import check_random_state
+
+
+def set_random_seeds(seed, cuda):
+    """Set seeds for python random module numpy.random and torch.
+
+    Parameters
+    ----------
+    seed: int
+        Random seed.
+    cuda: bool
+        Whether to set cuda seed with torch.
+    """
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if cuda:
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+
+
+def np_to_var(
+    X, requires_grad=False, dtype=None, pin_memory=False, **tensor_kwargs
+):
+    """
+    Convenience function to transform numpy array to `torch.Tensor`.
+
+    Converts `X` to ndarray using asarray if necessary.
+
+    Parameters
+    ----------
+    X: ndarray or list or number
+        Input arrays
+    requires_grad: bool
+        passed on to Variable constructor
+    dtype: numpy dtype, optional
+    var_kwargs:
+        passed on to Variable constructor
+
+    Returns
+    -------
+    var: `torch.Tensor`
+    """
+    if not hasattr(X, "__len__"):
+        X = [X]
+    X = np.asarray(X)
+    if dtype is not None:
+        X = X.astype(dtype)
+    X_tensor = torch.tensor(X, requires_grad=requires_grad, **tensor_kwargs)
+    if pin_memory:
+        X_tensor = X_tensor.pin_memory()
+    return X_tensor
+
+
+def var_to_np(var):
+    """Convenience function to transform `torch.Tensor` to numpy
+    array.
+
+    Should work both for CPU and GPU."""
+    return var.cpu().data.numpy()
+
+
+def corr(a, b):
+    """
+    Computes correlation only between terms of a and terms of b, not within
+    a and b.
+
+    Parameters
+    ----------
+    a, b: 2darray, features x samples
+
+    Returns
+    -------
+    Correlation between features in x and features in y
+    """
+    # Difference to numpy:
+    # Correlation only between terms of x and y
+    # not between x and x or y and y
+    this_cov = cov(a, b)
+    return _cov_to_corr(this_cov, a, b)
+
+
+def cov(a, b):
+    """
+    Computes covariance only between terms of a and terms of b, not within
+    a and b.
+
+    Parameters
+    ----------
+    a, b: 2darray, features x samples
+
+    Returns
+    -------
+    Covariance between features in x and features in y
+    """
+    demeaned_a = a - np.mean(a, axis=1, keepdims=True)
+    demeaned_b = b - np.mean(b, axis=1, keepdims=True)
+    this_cov = np.dot(demeaned_a, demeaned_b.T) / (b.shape[1] - 1)
+    return this_cov
+
+
+def _cov_to_corr(this_cov, a, b):
+    # computing "unbiased" corr
+    # ddof=1 for unbiased..
+    var_a = np.var(a, axis=1, ddof=1)
+    var_b = np.var(b, axis=1, ddof=1)
+    return _cov_and_var_to_corr(this_cov, var_a, var_b)
+
+
+def _cov_and_var_to_corr(this_cov, var_a, var_b):
+    divisor = np.outer(np.sqrt(var_a), np.sqrt(var_b))
+    return this_cov / divisor
+
+
+def wrap_reshape_apply_fn(stat_fn, a, b, axis_a, axis_b):
+    """
+    Reshape two nd-arrays into 2d-arrays, apply function and reshape
+    result back.
+
+    Parameters
+    ----------
+    stat_fn: function
+        Function to apply to 2d-arrays
+    a: nd-array: nd-array
+    b: nd-array
+    axis_a: int or list of int
+        sample axis
+    axis_b: int or list of int
+        sample axis
+
+    Returns
+    -------
+    result: nd-array
+        The result reshaped to remaining_dims_a + remaining_dims_b
+    """
+    if not hasattr(axis_a, "__len__"):
+        axis_a = [axis_a]
+    if not hasattr(axis_b, "__len__"):
+        axis_b = [axis_b]
+    other_axis_a = [i for i in range(a.ndim) if i not in axis_a]
+    other_axis_b = [i for i in range(b.ndim) if i not in axis_b]
+    transposed_topo_a = a.transpose(tuple(other_axis_a) + tuple(axis_a))
+    n_stat_axis_a = [a.shape[i] for i in axis_a]
+    n_other_axis_a = [a.shape[i] for i in other_axis_a]
+    flat_topo_a = transposed_topo_a.reshape(
+        np.prod(n_other_axis_a), np.prod(n_stat_axis_a)
+    )
+    transposed_topo_b = b.transpose(tuple(other_axis_b) + tuple(axis_b))
+    n_stat_axis_b = [b.shape[i] for i in axis_b]
+    n_other_axis_b = [b.shape[i] for i in other_axis_b]
+    flat_topo_b = transposed_topo_b.reshape(
+        np.prod(n_other_axis_b), np.prod(n_stat_axis_b)
+    )
+    assert np.array_equal(n_stat_axis_a, n_stat_axis_b)
+    stat_result = stat_fn(flat_topo_a, flat_topo_b)
+    topo_result = stat_result.reshape(
+        tuple(n_other_axis_a) + tuple(n_other_axis_b)
+    )
+    return topo_result
+
+
+def get_balanced_batches(
+    n_trials, rng, shuffle, n_batches=None, batch_size=None
+):
+    """Create indices for batches balanced in size
+    (batches will have maximum size difference of 1).
+    Supply either batch size or number of batches. Resulting batches
+    will not have the given batch size but rather the next largest batch size
+    that allows to split the set into balanced batches (maximum size difference 1).
+
+    Parameters
+    ----------
+    n_trials : int
+        Size of set.
+    rng : RandomState
+    shuffle : bool
+        Whether to shuffle indices before splitting set.
+    n_batches : int, optional
+    batch_size : int, optional
+
+    Returns
+    -------
+    batches: list of list of int
+        Indices for each batch.
+    """
+    assert batch_size is not None or n_batches is not None
+    if n_batches is None:
+        n_batches = int(np.round(n_trials / float(batch_size)))
+
+    if n_batches > 0:
+        min_batch_size = n_trials // n_batches
+        n_batches_with_extra_trial = n_trials % n_batches
+    else:
+        n_batches = 1
+        min_batch_size = n_trials
+        n_batches_with_extra_trial = 0
+    assert n_batches_with_extra_trial < n_batches
+    all_inds = np.array(range(n_trials))
+    if shuffle:
+        rng.shuffle(all_inds)
+    i_start_trial = 0
+    i_stop_trial = 0
+    batches = []
+    for i_batch in range(n_batches):
+        i_stop_trial += min_batch_size
+        if i_batch < n_batches_with_extra_trial:
+            i_stop_trial += 1
+        batch_inds = all_inds[range(i_start_trial, i_stop_trial)]
+        batches.append(batch_inds)
+        i_start_trial = i_stop_trial
+    assert i_start_trial == n_trials
+    return batches
+
+
+def create_mne_dummy_raw(n_channels, n_times, sfreq, include_anns=True,
+                         description=None, savedir=None, save_format='fif',
+                         overwrite=True, random_state=None):
+    """Create an mne.io.RawArray with fake data, and optionally save it.
+
+    This will overwrite already existing files.
+
+    Parameters
+    ----------
+    n_channels : int
+        Number of channels.
+    n_times : int
+        Number of samples.
+    sfreq : float
+        Sampling frequency.
+    include_anns : bool
+        If True, also create annotations.
+    description : list | None
+        List of descriptions used for creating annotations. It should contain
+        10 elements.
+    savedir : str | None
+        If provided as a string, the file will be saved under that directory.
+    save_format : str | list
+        If `savedir` is provided, this specifies the file format the data should
+        be saved to. Can be 'raw' or 'hdf5', or a list containing both.
+    random_state : int | RandomState
+        Random state for the generation of random data.
+
+    Returns
+    -------
+    raw : mne.io.Raw
+        The created Raw object.
+    save_fname : dict | None
+        Dictionary containing the name the raw data was saved to.
+    """
+    random_state = check_random_state(random_state)
+    data = random_state.rand(n_channels, n_times)
+    ch_names = [f'ch{i}' for i in range(n_channels)]
+    ch_types = ['eeg'] * n_channels
+    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
+
+    raw = mne.io.RawArray(data, info)
+
+    if include_anns:
+        n_anns = 10
+        inds = np.linspace(
+            int(sfreq * 2), int(n_times - sfreq * 2), num=n_anns).astype(int)
+        onset = raw.times[inds]
+        duration = [1] * n_anns
+        if description is None:
+            description = ['test'] * n_anns
+        anns = mne.Annotations(onset, duration, description)
+        raw = raw.set_annotations(anns)
+
+    save_fname = dict()
+    if savedir is not None:
+        if not isinstance(save_format, list):
+            save_format = [save_format]
+        fname = os.path.join(savedir, 'fake_eeg_raw')
+
+        if 'fif' in save_format:
+            fif_fname = fname + '.fif'
+            raw.save(fif_fname, overwrite=overwrite)
+            save_fname['fif'] = fif_fname
+        if 'hdf5' in save_format:
+            h5_fname = fname + '.h5'
+            with h5py.File(h5_fname, 'w') as f:
+                f.create_dataset(
+                    'fake_raw', dtype='f8', data=raw.get_data())
+            save_fname['hdf5'] = h5_fname
+
+    return raw, save_fname
+
+
+class ThrowAwayIndexLoader(object):
+    def __init__(self, net, loader, is_regression):
+        self.net = net
+        self.loader = loader
+        self.last_i = None
+        self.is_regression = is_regression
+
+    def __iter__(self, ):
+        normal_iter = self.loader.__iter__()
+        for batch in normal_iter:
+            if len(batch) == 3:
+                x, y, i = batch
+                # Store for scoring callbacks
+                self.net._last_window_inds = i
+            else:
+                x, y = batch
+
+            # TODO: should be on dataset side
+            if hasattr(x, 'type'):
+                x = x.type(torch.float32)
+                if self.is_regression:
+                    y = y.type(torch.float32)
+                else:
+                    y = y.type(torch.int64)
+            yield x, y
+
+
+def update_estimator_docstring(base_class, docstring):
+    base_doc = base_class.__doc__.replace(' : ', ': ')
+    idx = base_doc.find('callbacks:')
+    idx_end = idx + base_doc[idx:].find('\n\n')
+    # remove callback descripiton already included in braindecode docstring
+    filtered_doc = base_doc[:idx] + base_doc[idx_end+6:]
+    splitted = docstring.split('Parameters\n    ----------\n    ')
+    out_docstring = splitted[0] + \
+                    filtered_doc[filtered_doc.find('Parameters'):filtered_doc.find('Attributes')] + \
+                    splitted[1] + \
+                    filtered_doc[filtered_doc.find('Attributes'):]
+    return out_docstring
+
+
+def read_all_file_names(directory, extension):
+    """Read all files with specified extension from given path and sorts them
+    based on a given sorting key.
+
+    Parameters
+    ----------
+    directory: str
+        Parent directory to be searched for files of the specified type.
+    extension: str
+        File extension, i.e. ".edf" or ".txt".
+
+    Returns
+    -------
+    file_paths: list(str)
+        List of all files found in (sub)directories of path.
+    """
+    assert extension.startswith('.')
+    file_paths = glob.glob(directory + '**/*' + extension, recursive=True)
+    assert len(file_paths) > 0, (
+        f'something went wrong. Found no {extension} files in {directory}')
+    return file_paths

braindecode/version.py

@@ -0,0 +1 @@
+__version__ = "0.5.1"

braindecode/visualization/__init__.py

@@ -0,0 +1,5 @@
+"""
+Functions for visualisations, especially of the ConvNets.
+"""
+
+from .gradients import compute_amplitude_gradients

braindecode/visualization/gradients.py

@@ -0,0 +1,53 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+
+import numpy as np
+from skorch.utils import to_numpy, to_tensor
+import torch
+
+
+def compute_amplitude_gradients(model, dataset, batch_size):
+    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
+                                         drop_last=False, shuffle=False)
+    all_amp_grads = []
+    for batch_X, _, _ in loader:
+        this_amp_grads = compute_amplitude_gradients_for_X(model, batch_X, )
+        all_amp_grads.append(this_amp_grads)
+    all_amp_grads = np.concatenate(all_amp_grads, axis=1)
+    return all_amp_grads
+
+
+def compute_amplitude_gradients_for_X(model, X):
+    device = next(model.parameters()).device
+    ffted = np.fft.rfft(X, axis=2)
+    amps = np.abs(ffted)
+    phases = np.angle(ffted)
+    amps_th = to_tensor(amps.astype(np.float32), device=device).requires_grad_(True)
+    phases_th = to_tensor(phases.astype(np.float32), device=device).requires_grad_(True)
+
+    fft_coefs = amps_th.unsqueeze(-1) * torch.stack(
+        (torch.cos(phases_th), torch.sin(phases_th)), dim=-1)
+    fft_coefs = fft_coefs.squeeze(3)
+
+    try:
+        complex_fft_coefs = torch.view_as_complex(fft_coefs)
+        iffted = torch.fft.irfft(
+            complex_fft_coefs, n=X.shape[2], dim=2)
+    except AttributeError:
+        iffted = torch.irfft(  # Deprecated since 1.7
+            fft_coefs, signal_ndim=1, signal_sizes=(X.shape[2],))
+
+    outs = model(iffted)
+
+    n_filters = outs.shape[1]
+    amp_grads_per_filter = np.full((n_filters,) + ffted.shape,
+                                   np.nan, dtype=np.float32)
+    for i_filter in range(n_filters):
+        mean_out = torch.mean(outs[:, i_filter])
+        mean_out.backward(retain_graph=True)
+        amp_grads = to_numpy(amps_th.grad.clone())
+        amp_grads_per_filter[i_filter] = amp_grads
+        amps_th.grad.zero_()
+    assert not np.any(np.isnan(amp_grads_per_filter))
+    return amp_grads_per_filter

test/acceptance_tests/test_cropped_decoding.py

@@ -0,0 +1,171 @@
+# Authors: Maciej Sliwowski
+#          Robin Tibor Schirrmeister
+#
+# License: BSD-3
+
+import mne
+import numpy as np
+import torch
+from mne.io import concatenate_raws
+from skorch.helper import predefined_split
+from torch import optim
+
+from braindecode import EEGClassifier
+from braindecode.datautil.xy import create_from_X_y
+from braindecode.training.losses import CroppedLoss
+from braindecode.models import ShallowFBCSPNet
+from braindecode.models.util import to_dense_prediction_model, get_output_shape
+from braindecode.util import set_random_seeds
+
+
+def test_cropped_decoding():
+    # 5,6,7,10,13,14 are codes for executed and imagined hands/feet
+    subject_id = 1
+    event_codes = [5, 6, 9, 10, 13, 14]
+
+    # This will download the files if you don't have them yet,
+    # and then return the paths to the files.
+    physionet_paths = mne.datasets.eegbci.load_data(
+        subject_id, event_codes, update_path=False
+    )
+
+    # Load each of the files
+    parts = [
+        mne.io.read_raw_edf(
+            path, preload=True, stim_channel="auto", verbose="WARNING"
+        )
+        for path in physionet_paths
+    ]
+
+    # Concatenate them
+    raw = concatenate_raws(parts)
+
+    # Find the events in this dataset
+    events, _ = mne.events_from_annotations(raw)
+    # Use only EEG channels
+    eeg_channel_inds = mne.pick_types(
+        raw.info, meg=False, eeg=True, stim=False, eog=False, exclude="bads"
+    )
+
+    # Extract trials, only using EEG channels
+    epoched = mne.Epochs(
+        raw,
+        events,
+        dict(hands=2, feet=3),
+        tmin=1,
+        tmax=4.1,
+        proj=False,
+        picks=eeg_channel_inds,
+        baseline=None,
+        preload=True,
+    )
+    # Convert data from volt to millivolt
+    # Pytorch expects float32 for input and int64 for labels.
+    X = (epoched.get_data() * 1e6).astype(np.float32)
+    y = (epoched.events[:, 2] - 2).astype(np.int64)  # 2,3 -> 0,1
+
+    # Set if you want to use GPU
+    # You can also use torch.cuda.is_available() to determine if cuda is available on your machine.
+    cuda = False
+    set_random_seeds(seed=20170629, cuda=cuda)
+
+    # This will determine how many crops are processed in parallel
+    input_window_samples = 450
+    n_classes = 2
+    in_chans = X.shape[1]
+    # final_conv_length determines the size of the receptive field of the ConvNet
+    model = ShallowFBCSPNet(
+        in_chans=in_chans,
+        n_classes=n_classes,
+        input_window_samples=input_window_samples,
+        final_conv_length=12,
+    )
+    to_dense_prediction_model(model)
+
+    if cuda:
+        model.cuda()
+
+    # Perform forward pass to determine how many outputs per input
+    n_preds_per_input = get_output_shape(model, in_chans, input_window_samples)[2]
+
+    train_set = create_from_X_y(X[:60], y[:60],
+                                drop_last_window=False,
+                                window_size_samples=input_window_samples,
+                                window_stride_samples=n_preds_per_input)
+
+    valid_set = create_from_X_y(X[60:], y[60:],
+                                drop_last_window=False,
+                                window_size_samples=input_window_samples,
+                                window_stride_samples=n_preds_per_input)
+
+
+    train_split = predefined_split(valid_set)
+
+    clf = EEGClassifier(
+        model,
+        cropped=True,
+        criterion=CroppedLoss,
+        criterion__loss_function=torch.nn.functional.nll_loss,
+        optimizer=optim.Adam,
+        train_split=train_split,
+        batch_size=32,
+        callbacks=['accuracy'],
+    )
+
+    clf.fit(train_set, y=None, epochs=4)
+
+    np.testing.assert_allclose(
+        clf.history[:, 'train_loss'],
+        np.array(
+            [
+                1.6666231592496237,
+                1.2292670885721841,
+                1.1270817518234253,
+                1.1752660751342774
+            ]
+        ),
+        rtol=1e-3,
+        atol=1e-4,
+    )
+
+    np.testing.assert_allclose(
+        clf.history[:, 'valid_loss'],
+        np.array(
+            [
+                1.5687058925628663,
+                0.8510023872057597,
+                2.087181798617045,
+                0.7100235184033712
+            ]
+        ),
+        rtol=1e-3,
+        atol=1e-3,
+    )
+
+    np.testing.assert_allclose(
+        clf.history[:, 'train_accuracy'],
+        np.array(
+            [
+                0.48333333333333334,
+                0.5,
+                0.5,
+                0.6333333333333333
+            ]
+        ),
+        rtol=1e-3,
+        atol=1e-4,
+    )
+
+    np.testing.assert_allclose(
+        clf.history[:, 'valid_accuracy'],
+        np.array(
+            [
+                0.533333,
+                0.5,
+                0.466667,
+                0.666667
+            ]
+        ),
+        rtol=1e-3,
+        atol=1e-4,
+    )