MRG, ENH: Add mne.export

doc/changes/latest.inc

@@ -39,6 +39,8 @@ Enhancements
 
 - New function :func:`mne.label.find_pos_in_annot` to get atlas label for MRI coordinates.  (:gh:`9376` by **by new contributor** |Marian Dovgialo|_)
 
+- New namespace `mne.export` created to contain functions (such as `mne.export.export_raw` and `mne.export.export_epochs`) for exporting data to non-FIF formats (:gh:`9427` by `Eric Larson`_)
+
 - Add support for Hitachi fNIRS devices in `mne.io.read_raw_hitachi` (:gh:`9391` by `Eric Larson`_)
 
 - Add support for ``picks`` in :func:`mne.stc_near_sensors` (:gh:`9396` by `Eric Larson`_)

doc/export.rst

@@ -0,0 +1,17 @@
+
+Exporting
+================
+
+:py:mod:`mne.export`:
+
+.. automodule:: mne.export
+   :no-members:
+   :no-inherited-members:
+
+.. currentmodule:: mne.export
+
+.. autosummary::
+   :toctree: generated/
+
+   export_epochs
+   export_raw

doc/python_reference.rst

@@ -27,6 +27,7 @@ directly from a terminal, see :ref:`python_commands`.
     reading_raw_data
     file_io
     creating_from_arrays
+    export
     datasets
     visualization
     preprocessing

mne/__init__.py

@@ -121,6 +121,7 @@
 from . import time_frequency
 from . import viz
 from . import decoding
+from . import export
 
 # initialize logging
 set_log_level(None, False)

mne/conftest.py

@@ -38,6 +38,9 @@
 fname_trans = op.join(s_path, 'sample_audvis_trunc-trans.fif')
 
 
+collect_ignore = ['export/_eeglab.py']
+
+
 def pytest_configure(config):
     """Configure pytest options."""
     # Markers

mne/epochs.py

@@ -20,7 +20,6 @@
 
 import numpy as np
 
-from .io.utils import _get_als_coords_from_chs
 from .io.write import (start_file, start_block, end_file, end_block,
                        write_int, write_float, write_float_matrix,
                        write_double_matrix, write_complex_float_matrix,
@@ -52,15 +51,14 @@
 from .utils import (_check_fname, check_fname, logger, verbose,
                     _time_mask, check_random_state, warn, _pl,
                     sizeof_fmt, SizeMixin, copy_function_doc_to_method_doc,
-                    _check_pandas_installed, _check_eeglabio_installed,
+                    _check_pandas_installed,
                     _check_preload, GetEpochsMixin,
                     _prepare_read_metadata, _prepare_write_metadata,
                     _check_event_id, _gen_events, _check_option,
                     _check_combine, ShiftTimeMixin, _build_data_frame,
                     _check_pandas_index_arguments, _convert_times,
                     _scale_dataframe_data, _check_time_format, object_size,
-                    _on_missing, _validate_type, _ensure_events,
-                    _infer_check_export_fmt)
+                    _on_missing, _validate_type, _ensure_events)
 from .utils.docs import fill_doc
 from .data.html_templates import epochs_template
 
@@ -1836,37 +1834,8 @@ def export(self, fname, fmt='auto', verbose=None):
         -----
         %(export_eeglab_note)s
         """
-        supported_export_formats = {
-            'eeglab': ('set',),
-            'edf': ('edf',),
-            'brainvision': ('eeg', 'vmrk', 'vhdr',)
-        }
-        fmt = _infer_check_export_fmt(fmt, fname, supported_export_formats)
-
-        if fmt == 'eeglab':
-            _check_eeglabio_installed()
-            import eeglabio.epochs
-            # load data first
-            self.load_data()
-
-            # remove extra epoc and STI channels
-            drop_chs = ['epoc', 'STI 014']
-            ch_names = [ch for ch in self.ch_names if ch not in drop_chs]
-            cart_coords = _get_als_coords_from_chs(self.info['chs'],
-                                                   drop_chs)
-
-            eeglabio.epochs.export_set(fname,
-                                       data=self.get_data(picks=ch_names),
-                                       sfreq=self.info['sfreq'],
-                                       events=self.events,
-                                       tmin=self.tmin, tmax=self.tmax,
-                                       ch_names=ch_names,
-                                       event_id=self.event_id,
-                                       ch_locs=cart_coords)
-        elif fmt == 'edf':
-            raise NotImplementedError('Export to EDF format not implemented.')
-        elif fmt == 'brainvision':
-            raise NotImplementedError('Export to BrainVision not implemented.')
+        from .export import export_epochs
+        export_epochs(fname, self, fmt, verbose)
 
     def equalize_event_counts(self, event_ids=None, method='mintime'):
         """Equalize the number of trials in each condition.

mne/export/__init__.py

@@ -0,0 +1 @@
+from ._export import export_raw, export_epochs

mne/export/_eeglab.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+# Authors: MNE Developers
+#
+# License: BSD (3-clause)
+
+import numpy as np
+
+from ..utils import _check_eeglabio_installed
+_check_eeglabio_installed()
+import eeglabio.raw  # noqa: E402
+import eeglabio.epochs  # noqa: E402
+
+
+def _export_raw(fname, raw):
+    # load data first
+    raw.load_data()
+
+    # remove extra epoc and STI channels
+    drop_chs = ['epoc']
+    if not (raw.filenames[0].endswith('.fif')):
+        drop_chs.append('STI 014')
+
+    ch_names = [ch for ch in raw.ch_names if ch not in drop_chs]
+    cart_coords = _get_als_coords_from_chs(raw.info['chs'], drop_chs)
+
+    annotations = [raw.annotations.description,
+                   raw.annotations.onset,
+                   raw.annotations.duration]
+    eeglabio.raw.export_set(
+        fname, data=raw.get_data(picks=ch_names), sfreq=raw.info['sfreq'],
+        ch_names=ch_names, ch_locs=cart_coords, annotations=annotations)
+
+
+def _export_epochs(fname, epochs):
+    _check_eeglabio_installed()
+    # load data first
+    epochs.load_data()
+
+    # remove extra epoc and STI channels
+    drop_chs = ['epoc', 'STI 014']
+    ch_names = [ch for ch in epochs.ch_names if ch not in drop_chs]
+    cart_coords = _get_als_coords_from_chs(epochs.info['chs'], drop_chs)
+
+    eeglabio.epochs.export_set(
+        fname, data=epochs.get_data(picks=ch_names),
+        sfreq=epochs.info['sfreq'], events=epochs.events,
+        tmin=epochs.tmin, tmax=epochs.tmax, ch_names=ch_names,
+        event_id=epochs.event_id, ch_locs=cart_coords)
+
+
+def _get_als_coords_from_chs(chs, drop_chs=None):
+    """Extract channel locations in ALS format (x, y, z) from a chs instance.
+
+    Returns
+    -------
+    None if no valid coordinates are found (all zeros)
+    """
+    if drop_chs is None:
+        drop_chs = []
+    cart_coords = np.array([d['loc'][:3] for d in chs
+                            if d['ch_name'] not in drop_chs])
+    if cart_coords.any():  # has coordinates
+        # (-y x z) to (x y z)
+        cart_coords[:, 0] = -cart_coords[:, 0]  # -y to y
+        # swap x (1) and y (0)
+        cart_coords[:, [0, 1]] = cart_coords[:, [1, 0]]
+    else:
+        cart_coords = None
+    return cart_coords

mne/export/_export.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+# Authors: MNE Developers
+#
+# License: BSD (3-clause)
+
+import os.path as op
+
+from ..utils import verbose, _validate_type
+
+
+@verbose
+def export_raw(fname, raw, fmt='auto', verbose=None):
+    """Export Raw to external formats.
+
+    Supported formats: EEGLAB (set, uses :mod:`eeglabio`)
+    %(export_warning)s
+
+    Parameters
+    ----------
+    %(export_params_fname)s
+    raw : instance of Raw
+        The raw instance to export.
+    %(export_params_fmt)s
+    %(verbose)s
+
+    Notes
+    -----
+    %(export_eeglab_note)s
+    """
+    supported_export_formats = {  # format : extensions
+        'eeglab': ('set',),
+        'edf': ('edf',),
+        'brainvision': ('eeg', 'vmrk', 'vhdr',)
+    }
+    fmt = _infer_check_export_fmt(fmt, fname, supported_export_formats)
+
+    if fmt == 'eeglab':
+        from ._eeglab import _export_raw
+        _export_raw(fname, raw)
+    elif fmt == 'edf':
+        raise NotImplementedError('Export to EDF format not implemented.')
+    elif fmt == 'brainvision':
+        raise NotImplementedError('Export to BrainVision not implemented.')
+
+
+@verbose
+def export_epochs(fname, epochs, fmt='auto', verbose=None):
+    """Export Epochs to external formats.
+
+    Supported formats: EEGLAB (set, uses :mod:`eeglabio`)
+    %(export_warning)s
+
+    Parameters
+    ----------
+    %(export_params_fname)s
+    epochs : instance of Epochs
+        The epochs to export.
+    %(export_params_fmt)s
+    %(verbose)s
+
+    Notes
+    -----
+    %(export_eeglab_note)s
+    """
+    supported_export_formats = {
+        'eeglab': ('set',),
+        'edf': ('edf',),
+        'brainvision': ('eeg', 'vmrk', 'vhdr',)
+    }
+    fmt = _infer_check_export_fmt(fmt, fname, supported_export_formats)
+
+    if fmt == 'eeglab':
+        from ._eeglab import _export_epochs
+        _export_epochs(fname, epochs)
+    elif fmt == 'edf':
+        raise NotImplementedError('Export to EDF format not implemented.')
+    elif fmt == 'brainvision':
+        raise NotImplementedError('Export to BrainVision not implemented.')
+
+
+def _infer_check_export_fmt(fmt, fname, supported_formats):
+    """Infer export format from filename extension if auto.
+
+    Raises error if fmt is auto and no file extension found,
+    then checks format against supported formats, raises error if format is not
+    supported.
+
+    Parameters
+    ----------
+    fmt : str
+        Format of the export, will only infer the format from filename if fmt
+        is auto.
+    fname : str
+        Name of the target export file, only used when fmt is auto.
+    supported_formats : dict of str : tuple/list
+        Dictionary containing supported formats (as keys) and each format's
+        corresponding file extensions in a tuple/list (e.g. 'eeglab': ('set',))
+    """
+    _validate_type(fmt, str, 'fmt')
+    fmt = fmt.lower()
+    if fmt == "auto":
+        fmt = op.splitext(fname)[1]
+        if fmt:
+            fmt = fmt[1:].lower()
+            # find fmt in supported formats dict's tuples
+            fmt = next((k for k, v in supported_formats.items() if fmt in v),
+                       fmt)  # default to original fmt for raising error later
+        else:
+            raise ValueError(f"Couldn't infer format from filename {fname}"
+                             " (no extension found)")
+
+    if fmt not in supported_formats:
+        supported = []
+        for format, extensions in supported_formats.items():
+            ext_str = ', '.join(f'*.{ext}' for ext in extensions)
+            supported.append(f'{format} ({ext_str})')
+
+        supported_str = ', '.join(supported)
+        raise ValueError(f"Format '{fmt}' is not supported. "
+                         f"Supported formats are {supported_str}.")
+    return fmt

mne/export/tests/test_export.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+"""Test exporting functions."""
+# Authors: MNE Developers
+#
+# License: BSD (3-clause)
+
+from pathlib import Path
+import os.path as op
+
+import pytest
+import numpy as np
+from numpy.testing import assert_allclose, assert_array_equal
+
+from mne import read_epochs_eeglab, Epochs
+from mne.tests.test_epochs import _get_data
+from mne.io import read_raw_fif, read_raw_eeglab
+from mne.utils import _check_eeglabio_installed
+
+
+@pytest.mark.skipif(not _check_eeglabio_installed(strict=False),
+                    reason='eeglabio not installed')
+def test_export_raw_eeglab(tmpdir):
+    """Test saving a Raw instance to EEGLAB's set format."""
+    fname = (Path(__file__).parent.parent.parent /
+             "io" / "tests" / "data" / "test_raw.fif")
+    raw = read_raw_fif(fname)
+    raw.load_data()
+    temp_fname = op.join(str(tmpdir), 'test.set')
+    raw.export(temp_fname)
+    raw.drop_channels([ch for ch in ['epoc']
+                       if ch in raw.ch_names])
+    raw_read = read_raw_eeglab(temp_fname, preload=True)
+    assert raw.ch_names == raw_read.ch_names
+    cart_coords = np.array([d['loc'][:3] for d in raw.info['chs']])  # just xyz
+    cart_coords_read = np.array([d['loc'][:3] for d in raw_read.info['chs']])
+    assert_allclose(cart_coords, cart_coords_read)
+    assert_allclose(raw.times, raw_read.times)
+    assert_allclose(raw.get_data(), raw_read.get_data())
+
+
+@pytest.mark.skipif(not _check_eeglabio_installed(strict=False),
+                    reason='eeglabio not installed')
+@pytest.mark.parametrize('preload', (True, False))
+def test_export_epochs_eeglab(tmpdir, preload):
+    """Test saving an Epochs instance to EEGLAB's set format."""
+    raw, events = _get_data()[:2]
+    raw.load_data()
+    epochs = Epochs(raw, events, preload=preload)
+    temp_fname = op.join(str(tmpdir), 'test.set')
+    epochs.export(temp_fname)
+    epochs.drop_channels([ch for ch in ['epoc', 'STI 014']
+                          if ch in epochs.ch_names])
+    epochs_read = read_epochs_eeglab(temp_fname)
+    assert epochs.ch_names == epochs_read.ch_names
+    cart_coords = np.array([d['loc'][:3]
+                           for d in epochs.info['chs']])  # just xyz
+    cart_coords_read = np.array([d['loc'][:3]
+                                for d in epochs_read.info['chs']])
+    assert_allclose(cart_coords, cart_coords_read)
+    assert_array_equal(epochs.events[:, 0],
+                       epochs_read.events[:, 0])  # latency
+    assert epochs.event_id.keys() == epochs_read.event_id.keys()  # just keys
+    assert_allclose(epochs.times, epochs_read.times)
+    assert_allclose(epochs.get_data(), epochs_read.get_data())