Preserve order of elec_names (#12)

* Preserve order of elec_names

* Remove default scatter color

* Fix tests

* Add changelog entry

* Extend plotting examples

* Some style fixes

* Shorten lines to avoid scrollbars

* Change thumbnail

Co-authored-by: Stefan Appelhoff <[email protected]>

Co-authored-by: Stefan Appelhoff <[email protected]>

Author: cbrnr

docs/changes.rst

@@ -10,6 +10,8 @@ Changelog
 2.1.0 (unreleased)
 ------------------
 - Add a ``show_axis`` parameter to :func:`eeg_positions.plot_coords`, by `Clemens Brunner`_ (:github:`#7`)
+- Add a ``sort`` parameter to :func:`eeg_positions.get_coords`, by `Clemens Brunner`_ (:github:`#12`)
+- Allow passing a list of colors to :func:`eeg_positions.plot_coords`, by `Clemens Brunner`_ (:github:`#12`)
 
 2.0.0 (2021-02-13)
 ------------------

eeg_positions/compute.py

@@ -152,6 +152,7 @@ def get_elec_coords(
     dim="2d",
     as_mne_montage=False,
     equator="Nz-T10-Iz-T9",
+    sort=False,
 ):
     """Get standard EEG electrode coordinates.
 
@@ -213,6 +214,9 @@ def get_elec_coords(
         as the equator, several electrodes may be drawn outside a circular
         head shape when projecting to 2D.
         Defaults to ``"Nz-T10-Iz-T9"``.
+    sort : bool
+        Whether to sort the returned coordinates alphabetically. If ``False`` (default),
+        preserve the order of ``elec_names`` if available.
 
     Returns
     -------
@@ -423,9 +427,14 @@ def get_elec_coords(
     # --------------------
     if len(elec_names) > 0:
         selection = df.label.isin(elec_names)
+        df_selection = df.loc[selection, :].copy()
+        if not sort:
+            df_selection = (
+                df_selection.set_index("label").reindex(elec_names).reset_index()
+            )
     else:
         selection = df.label.isin(system + LANDMARKS)
-    df_selection = df.loc[selection, :].copy()
+        df_selection = df.loc[selection, :].copy()
 
     # add special elec positions
     pos_to_add = {}
@@ -516,8 +525,9 @@ def get_elec_coords(
         df_selection.loc[:, "y"] = ys
 
     df_selection = df_selection.drop_duplicates(subset=["label"])
-    coords = df_selection.sort_values(by="label")
-    coords = coords.reset_index(drop=True)
+    if sort:
+        df_selection = df_selection.sort_values(by="label")
+    coords = df_selection.reset_index(drop=True)
     return coords
 
 
@@ -555,6 +565,7 @@ def _produce_files_and_do_x(x="save"):
                     dim=dim.lower(),
                     as_mne_montage=False,
                     equator=equator,
+                    sort=True,
                 )
 
                 fname = fname_template.format(equator, system, dim)

eeg_positions/tests/test_compute.py

@@ -19,14 +19,16 @@
     ("2d", "3d"),
     (True, False),
     ("Nz-T10-Iz-T9", "Fpz-T8-Oz-T7"),
+    (True, False),
 )
 
 
 @pytest.mark.parametrize(
-    "system, elec_names, drop_landmarks, dim, as_mne_montage, equator", valid_inputs
+    "system, elec_names, drop_landmarks, dim, as_mne_montage, equator, sort",
+    valid_inputs,
 )
 def test_get_elec_coords(
-    system, elec_names, drop_landmarks, dim, as_mne_montage, equator
+    system, elec_names, drop_landmarks, dim, as_mne_montage, equator, sort
 ):
     """Smoke test the get_elec_coords function."""
     get_elec_coords(
@@ -36,6 +38,7 @@ def test_get_elec_coords(
         dim=dim,
         as_mne_montage=as_mne_montage,
         equator=equator,
+        sort=sort,
     )
 
 
@@ -73,6 +76,15 @@ def test_get_elec_coords_io():
     with pytest.raises(ValueError, match=match):
         get_elec_coords(elec_names=["M1", "TP9"])
 
+    # check coords order
+    elec_names = ["Fp1", "AFz"]
+    coords = get_elec_coords(elec_names=elec_names)  # default sort=False
+    assert coords.label.to_list() == elec_names
+    coords = get_elec_coords(elec_names=elec_names, sort=False)
+    assert coords.label.to_list() == elec_names
+    coords = get_elec_coords(elec_names=elec_names, sort=True)
+    assert coords.label.to_list() == sorted(elec_names)
+
     # Mock a too-low version of mne
     mock_mne = mock.MagicMock()
     mock_mne.__version__ = "0.19.0"

eeg_positions/viz.py

@@ -178,16 +178,15 @@ def plot_coords(coords, scatter_kwargs={}, text_kwargs={}):
     dim = "3d" if "z" in coords.columns else "2d"
 
     # update kwargs
-    scatter_settings = dict(color="r")
+    scatter_settings = dict()
     scatter_settings.update(scatter_kwargs)
     text_settings = dict(fontsize=6)
     text_settings.update(text_kwargs)
 
     if dim == "2d":
         fig, ax = _plot_2d_head(RADIUS_INNER_CONTOUR)
-
+        ax.scatter(coords["x"], coords["y"], zorder=2.5, **scatter_settings)
         for _, row in coords.iterrows():
-            ax.scatter(row["x"], row["y"], zorder=2.5, **scatter_settings)
             ax.text(row["x"], row["y"], row["label"], **text_settings)
 
     else:

examples/plot_positions.py

@@ -12,63 +12,104 @@
 from eeg_positions import get_elec_coords, plot_coords
 
 # %%
-# Get the electrode positions!
-# Let's start with the basic 10-20 system in 2 dimensions (``"x"`` and ``"y"``).
+# Let's start with the basic 10-20 system in two dimensions:
 
 coords = get_elec_coords(
     system="1020",
     dim="2d",
 )
 
-# `coords` is a pandas.DataFrame object
+# %%
+# This function returns a ``pandas.DataFrame`` object:
+
 coords.head()
 
 # %%
 # Now let's plot these coordinates.
 # We can supply some style arguments to :func:`eeg_positions.plot_coords` to control
-# the color of the scatter dots, and the text annotations.
+# the color of the electrodes and the text annotations.
 
 fig, ax = plot_coords(
-    coords, scatter_kwargs=dict(color="green"), text_kwargs=dict(fontsize=10)
+    coords, scatter_kwargs={"color": "g"}, text_kwargs={"fontsize": 10}
 )
 
-ax.axis("off")
 fig
 
 # %%
-# Notice in the above plot, that the "landmarks" are there: ``NAS``, ``LPA``,
-# and ``RPA``. We can drop these by passing the ``drop_landmarks=True`` to
-# :func:`get_elec_coords`.
+# Notice that the "landmarks" ``NAS``, ``LPA``, and ``RPA`` are included. We can drop
+# these by passing ``drop_landmarks=True`` to :func:`get_elec_coords`:
 
 coords = get_elec_coords(
     system="1020",
     drop_landmarks=True,
     dim="2d",
 )
 
-
 fig, ax = plot_coords(
-    coords, scatter_kwargs=dict(color="green"), text_kwargs=dict(fontsize=10)
+    coords, scatter_kwargs={"color": "g"}, text_kwargs={"fontsize": 10}
 )
 
-ax.axis("off")
 fig
 
 # %%
-# We can also plot in 3D. Let's pick a system with more electrodes now.
+# Often, we might have a list of electrode names that we would like to plot. For
+# example, let's assume we have the following 64 channel labels (based on the 10-05
+# system):
+
+chans = """Fp1 AF7 AF3 F1 F3 F5 F7 Fp2 AF8 AF4 F2 F4 F6 F8 FT7 FC5 FC3
+FC1 C1 C3 C5 T7 TP7 CP5 CP3 CP1 FT8 FC6 FC4 FC2 C2 C4 C6 T8 TP8 CP6 CP4
+CP2 P1 P3 P5 P7 P9 PO7 PO3 O1 P2 P4 P6 P8 P10 PO8 PO4 O2 Iz Oz POz Pz
+Fz AFz Fpz CPz Cz FCz""".split()
+
+# %%
+# Many experiments aggregate electrodes into regions of interest (ROIs), which we could
+# visualize with different colors. Let's get their coordinates first:
+
+coords = get_elec_coords(elec_names=chans)
+
+# %%
+# Now we specifiy individual colors using the ``scatter_kwargs``` argument. We create a
+# list of 64 colors corresponding to our 64 coordinates (in the original order as
+# provided by ``chans``):
+
+colors = (
+    ["salmon"] * 14
+    + ["skyblue"] * 24
+    + ["violet"] * 16
+    + ["lightgreen"] * 7
+    + ["khaki"] * 3
+)
+
+# sphinx_gallery_thumbnail_number = 3
+fig, ax = plot_coords(
+    coords,
+    scatter_kwargs={
+        "s": 150,  # electrode size
+        "color": colors,
+        "edgecolors": "black",  # black electrode outline
+        "linewidths": 0.5,  # thin outline
+    },
+    text_kwargs={
+        "ha": "center",  # center electrode label horizontally
+        "va": "center",  # center electrode label vertically
+        "fontsize": 5,  # smaller font size
+    },
+)
+
+# %%
+# We can also plot in 3D. Let's pick a system with more electrodes now:
 
 coords = get_elec_coords(
     system="1010",
     drop_landmarks=True,
     dim="3d",
 )
 
-
 fig, ax = plot_coords(coords, text_kwargs=dict(fontsize=7))
 
 fig
 
 # %%
 # When using these commands from an interactive Python session, try to set
-# the Ipython magic `%matplotlib qt`, which will allow you to freely view the 3d
-# plot and rotate the camera.
+# the IPython magic ``%matplotlib`` or ``%matplotlib qt``, which will allow you to
+# freely view the 3D plot and rotate the camera.