[WIP] Support exporting mne objects to xarray DataArray objects

environment.yml

@@ -55,3 +55,4 @@ dependencies:
 - eeglabio
 - edflib-python
 - pybv
+- xarray

mne/utils/__init__.py

@@ -52,7 +52,7 @@
                        assert_object_equal, assert_and_remove_boundary_annot,
                        _raw_annot, assert_dig_allclose, assert_meg_snr,
                        assert_snr, assert_stcs_equal, _click_ch_name,
-                       requires_openmeeg_mark)
+                       requires_openmeeg_mark, requires_xarray)
 from .numerics import (hashfunc, _compute_row_norms,
                        _reg_pinv, random_permutation, _reject_data_segments,
                        compute_corr, _get_inst_data, array_split_idx,

mne/utils/_testing.py

@@ -130,6 +130,7 @@ def requires_module(function, name, call=None):
 requires_pylsl = partial(requires_module, name='pylsl')
 requires_sklearn = partial(requires_module, name='sklearn')
 requires_mne = partial(requires_module, name='MNE-C', call=_mne_call)
+requires_xarray = partial(requires_module, name='xarray')
 
 
 def requires_mne_mark():

mne/utils/tests/test_xarray.py

@@ -0,0 +1,21 @@
+import numpy as np
+import mne
+from mne.utils import requires_xarray
+
+
+@requires_xarray
+def test_conversion_to_xarray():
+    """Test conversion of mne object to xarray DataArray."""
+    import xarray as xr
+    from mne.utils.xarray import to_xarray
+
+    info = mne.create_info(list('abcd'), sfreq=250)
+    data = np.random.rand(4, 350)
+    erp = mne.EvokedArray(data, info, tmin=-0.5)
+
+    erp_x = to_xarray(erp)
+    assert isinstance(erp_x, xr.DataArray)
+    assert erp_x.shape == (4, 350)
+    assert erp_x.dims == ('chan', 'time')
+    assert erp_x.coords['chan'].data.tolist() == erp.ch_names
+    assert (erp_x.coords['time'].data == erp.times).all()

mne/utils/xarray.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from .. import Epochs, Evoked
+from ..time_frequency import AverageTFR
+
+
+def to_xarray(mne_inst):
+    """Convert MNE object instance to xarray DataArray.
+
+    Parameters
+    ----------
+    mne_inst : Epochs | Evoked
+        The MNE object to convert.
+
+    Returns
+    -------
+    xarr : DataArray
+        The xarray object.
+    """
+    from xarray import DataArray
+    from mne.utils import _validate_type
+
+    _validate_type(mne_inst, (Epochs, Evoked, AverageTFR))
+
+    if isinstance(mne_inst, Epochs):
+        data = mne_inst.get_data()
+        dims = ('chan', 'trial', 'time')
+    elif isinstance(mne_inst, Evoked):
+        data = mne_inst.data
+        dims = ('chan', 'time')
+    else:
+        raise ValueError('MNE instance must be Epochs or Evoked.')
+
+    coords = dict(chan=mne_inst.ch_names)
+    if 'time' in dims:
+        coords['time'] = mne_inst.times
+    if 'trial' in dims:
+        coords['trial'] = np.arange(mne_inst.n_epochs)
+
+    return DataArray(data, dims=dims, coords=coords)

requirements.txt

@@ -43,3 +43,4 @@ mne-qt-browser
 darkdetect
 qdarkstyle
 threadpoolctl
+xarray

tools/github_actions_dependencies.sh

@@ -28,6 +28,8 @@ else
 	echo "nilearn and openmeeg"
 	pip install $STD_ARGS --pre https://github.com/nilearn/nilearn/zipball/main
 	pip install $STD_ARGS --pre --only-binary ":all:" -i "https://test.pypi.org/simple" openmeeg
+	echo "xarray"
+	pip install --progress-bar off --pre xarray
 	echo "VTK"
 	pip install $STD_ARGS --pre --only-binary ":all:" vtk
 	python -c "import vtk"