ENH: different color for each volume source space (mne-tools#9114)

* different kind of sources with different colours * fix color selection * test plot alignment mixed source space test * add an example to check with CI * make proper color format * update changelog Co-authored-by: Eric Larson <[email protected]>
rbruna · Mar 18, 2021 · 0ddde8c · 0ddde8c
1 parent d8fc038
commit 0ddde8c
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 21 deletions.
diff --git a/doc/changes/latest.inc b/doc/changes/latest.inc
@@ -47,6 +47,8 @@ Current (0.23.dev0)
 
 Enhancements
 ~~~~~~~~~~~~
+- Add different colors for each volume source space in :func:`mne.viz.plot_alignment` (:gh:`9043` **by new contributor** |Valerii Chirkov|_)
+
 - Add ``overlap`` parameter to :func:`mne.make_fixed_length_epochs` to allow creating overlapping fixed length epochs (:gh:`9096` **by new contributor** |Silvia Cotroneo|_)
 
 - Add :meth:`mne.Dipole.to_mni` for more convenient  dipole.pos to MNI conversion (:gh:`9043` **by new contributor** |Valerii Chirkov|_)

diff --git a/examples/inverse/plot_mixed_source_space_inverse.py b/examples/inverse/plot_mixed_source_space_inverse.py
@@ -85,7 +85,7 @@
 #
 #    >>> write_source_spaces(fname_mixed_src, src, overwrite=True)
 #
-# We can also export source positions to nifti file and visualize it again:
+# We can also export source positions to NIfTI file and visualize it again:
 
 nii_fname = op.join(bem_dir, '%s-mixed-src.nii' % subject)
 src.export_volume(nii_fname, mri_resolution=True, overwrite=True)
@@ -149,7 +149,7 @@
                            subjects_dir=subjects_dir)
 ###############################################################################
 # Plot the volume
-# ----------------
+# ---------------
 
 fig = stc.volume().plot(initial_time=initial_time, src=src,
                         subjects_dir=subjects_dir)

diff --git a/mne/viz/_3d.py b/mne/viz/_3d.py
@@ -26,7 +26,7 @@
 from ..io.constants import FIFF
 from ..io.meas_info import read_fiducials, create_info
 from ..source_space import (_ensure_src, _create_surf_spacing, _check_spacing,
-                            _read_mri_info, SourceSpaces)
+                            _read_mri_info, SourceSpaces, read_freesurfer_lut)
 
 from ..surface import (get_meg_helmet_surf, _read_mri_surface, _DistanceQuery,
                        transform_surface_to, _project_onto_surface,
@@ -634,12 +634,6 @@ def plot_alignment(info=None, trans=None, subject=None, subjects_dir=None,
         if src_subject is not None and subject != src_subject:
             raise ValueError('subject ("%s") did not match the subject name '
                              ' in src ("%s")' % (subject, src_subject))
-        src_rr = np.concatenate([s['rr'][s['inuse'].astype(bool)]
-                                 for s in src])
-        src_nn = np.concatenate([s['nn'][s['inuse'].astype(bool)]
-                                 for s in src])
-    else:
-        src_rr = src_nn = np.empty((0, 3))
 
     if fwd is not None:
         _validate_type(fwd, [Forward])
@@ -846,9 +840,6 @@ def plot_alignment(info=None, trans=None, subject=None, subjects_dir=None,
         surfs[key] = transform_surface_to(surfs[key], coord_frame,
                                           [mri_trans, head_trans], copy=True)
 
-    if src is not None:
-        src_rr, src_nn = _update_coord_frame(src[0], src_rr, src_nn,
-                                             mri_trans, head_trans)
     if fwd is not None:
         fwd_rr, fwd_nn = _update_coord_frame(fwd, fwd_rr, fwd_nn,
                                              mri_trans, head_trans)
@@ -1053,14 +1044,36 @@ def plot_alignment(info=None, trans=None, subject=None, subjects_dir=None,
         surf = dict(rr=meg_rrs, tris=meg_tris)
         renderer.surface(surface=surf, color=color,
                          opacity=alpha, backface_culling=True)
-    if len(src_rr) > 0:
-        renderer.quiver3d(
-            x=src_rr[:, 0], y=src_rr[:, 1], z=src_rr[:, 2],
-            u=src_nn[:, 0], v=src_nn[:, 1], w=src_nn[:, 2],
-            color=(1., 1., 0.), mode='cylinder', scale=3e-3,
-            opacity=0.75, glyph_height=0.25,
-            glyph_center=(0., 0., 0.), glyph_resolution=20,
-            backface_culling=True)
+
+    if src is not None:
+        atlas_ids, colors = read_freesurfer_lut()
+        for ss in src:
+            src_rr = ss['rr'][ss['inuse'].astype(bool)]
+            src_nn = ss['nn'][ss['inuse'].astype(bool)]
+
+            src_rr, src_nn = _update_coord_frame(src[0], src_rr, src_nn,
+                                                 mri_trans, head_trans)
+            # volume sources
+            if ss['type'] == 'vol':
+                if ss['seg_name'] in colors.keys():
+                    color = colors[ss['seg_name']][:3]
+                    color = tuple(i / 256. for i in color)
+                else:
+                    color = (1., 1., 0.)
+
+            # surface and discrete sources
+            else:
+                color = (1., 1., 0.)
+
+            if len(src_rr) > 0:
+                renderer.quiver3d(
+                    x=src_rr[:, 0], y=src_rr[:, 1], z=src_rr[:, 2],
+                    u=src_nn[:, 0], v=src_nn[:, 1], w=src_nn[:, 2],
+                    color=color, mode='cylinder', scale=3e-3,
+                    opacity=0.75, glyph_height=0.25,
+                    glyph_center=(0., 0., 0.), glyph_resolution=20,
+                    backface_culling=True)
+
     if fwd is not None:
         red = (1.0, 0.0, 0.0)
         green = (0.0, 1.0, 0.0)

diff --git a/mne/viz/tests/test_3d.py b/mne/viz/tests/test_3d.py
@@ -165,7 +165,7 @@ def test_plot_evoked_field(renderer):
 @pytest.mark.slowtest  # can be slow on OSX
 @testing.requires_testing_data
 @traits_test
-def test_plot_alignment(tmpdir, renderer):
+def test_plot_alignment(tmpdir, renderer, mixed_fwd_cov_evoked):
     """Test plotting of -trans.fif files and MEG sensor layouts."""
     # generate fiducials file for testing
     tempdir = str(tmpdir)
@@ -209,6 +209,13 @@ def test_plot_alignment(tmpdir, renderer):
                   src=sample_src)
     sample_src.plot(subjects_dir=subjects_dir, head=True, skull=True,
                     brain='white')
+    # mixed source space
+    mixed_src = mixed_fwd_cov_evoked[0]['src']
+    assert mixed_src.kind == 'mixed'
+    plot_alignment(info, meg=['helmet', 'sensors'], dig=True,
+                   coord_frame='head', trans=Path(trans_fname),
+                   subject='sample', mri_fiducials=fiducials_path,
+                   subjects_dir=subjects_dir, src=mixed_src)
     renderer.backend._close_all()
     # no-head version
     renderer.backend._close_all()