Expose colormap in plot_sensors_connectivity

doc/whats_new.rst

@@ -24,16 +24,17 @@ Enhancements
 ~~~~~~~~~~~~
 
 - Add the option to set the number of connections plotted in :func:`mne_connectivity.viz.plot_sensors_connectivity` by `Qianliang Li`_ (:pr:`133`).
+- Allow setting colormap via new parameter ``cmap`` in :func:`mne_connectivity.viz.plot_sensors_connectivity` by `Daniel McCloy`_ (:pr:`141`).
 
 Bug
 ~~~
 
-- 
+-
 
 API
 ~~~
 
-- 
+-
 
 Authors
 ~~~~~~~

mne_connectivity/conftest.py

@@ -33,6 +33,9 @@ def pytest_configure(config):
     always::ResourceWarning
     # pydarkstyle
     ignore:.*Setting theme='dark' is not yet supported.*:RuntimeWarning
+    # imageio-ffmpeg (still happening as of version 0.4.8):
+    ignore:pkg_resources is deprecated as an API:DeprecationWarning
+    ignore:Deprecated call to `pkg_resources.declare_namespace.*:DeprecationWarning
     """  # noqa: E501
     for warning_line in warning_lines.split('\n'):
         warning_line = warning_line.strip()

mne_connectivity/viz/_3d.py

@@ -23,7 +23,7 @@
 @fill_doc
 def plot_sensors_connectivity(info, con, picks=None,
                               cbar_label='Connectivity',
-                              n_con=20):
+                              n_con=20, cmap='RdBu'):
     """Visualize the sensor connectivity in 3D.
 
     Parameters
@@ -38,6 +38,10 @@ def plot_sensors_connectivity(info, con, picks=None,
         Label for the colorbar.
     n_con : int
         Number of strongest connections shown. By default 20.
+    cmap : str | instance of matplotlib.colors.Colormap
+        Colormap for coloring connections by strength. If :class:`str`, must
+        be a valid Matplotlib colormap (i.e. a valid key of
+        ``matplotlib.colormaps``). Default is ``"RdBu"``.
 
     Returns
     -------
@@ -92,7 +96,8 @@ def plot_sensors_connectivity(info, con, picks=None,
                              destination=np.c_[x2, y2, z2],
                              scalars=np.c_[val, val],
                              vmin=vmin, vmax=vmax,
-                             reverse_lut=True)
+                             reverse_lut=True,
+                             colormap=cmap)
 
     renderer.scalarbar(source=tube, title=cbar_label)
 

mne_connectivity/viz/tests/test_3d.py

@@ -2,6 +2,9 @@
 import numpy as np
 import pytest
 
+from matplotlib import colormaps
+from numpy.testing import assert_almost_equal
+
 import mne
 from mne.datasets import testing
 
@@ -27,7 +30,17 @@
         plot_sensors_connectivity(info='foo', con=con, picks=picks)
     with pytest.raises(ValueError, match='does not correspond to the size'):
         plot_sensors_connectivity(info=info, con=con[::2, ::2], picks=picks)
-
-    fig = plot_sensors_connectivity(info=info, con=con, picks=picks)
+    cmap = 'viridis'
+    fig = plot_sensors_connectivity(info=info, con=con, picks=picks, cmap=cmap)
+    # check colormap
+    cmap_from_mpl = np.array(colormaps[cmap].colors)
+    cmap_from_vtk = np.array(
+        fig.plotter.scalar_bar.GetLookupTable().GetTable()
+    )
+    # discard alpha channel and convert uint8 -> norm
+    cmap_from_vtk = cmap_from_vtk[:, :3] / 255
+    cmap_from_vtk = cmap_from_vtk[::-1]  # for some reason order is flipped
+    assert_almost_equal(cmap_from_mpl, cmap_from_vtk, decimal=2)
+    # check title
     title = list(fig.plotter.scalar_bars.values())[0].GetTitle()
     assert title == 'Connectivity'