Refactor (#9121)

Author: GuillaumeFavelier

mne/viz/_brain/_brain.py

@@ -445,9 +445,7 @@ def __init__(self, subject_id, hemi, surf, title=None,
             self.silhouette = True
         else:
             self.silhouette = silhouette
-        # for now only one color bar can be added
-        # since it is the same for all figures
-        self._colorbar_added = False
+        self._scalar_bar = None
         # for now only one time label can be added
         # since it is the same for all figures
         self._time_label_added = False
@@ -692,14 +690,12 @@ def _clean(self):
         # Qt LeaveEvent requires _Iren so we use _FakeIren instead of None
         # to resolve the ref to vtkGenericRenderWindowInteractor
         self.plotter._Iren = _FakeIren()
-        if getattr(self.plotter, 'scalar_bar', None) is not None:
-            self.plotter.scalar_bar = None
         if getattr(self.plotter, 'picker', None) is not None:
             self.plotter.picker = None
         # XXX end PyVista
         for key in ('plotter', 'window', 'dock', 'tool_bar', 'menu_bar',
                     'status_bar', 'interactor', 'mpl_canvas', 'time_actor',
-                    'picked_renderer', 'act_data_smooth',
+                    'picked_renderer', 'act_data_smooth', '_scalar_bar',
                     'actions', 'widgets', 'geo', '_data'):
             setattr(self, key, None)
 
@@ -824,12 +820,11 @@ def _configure_time_label(self):
             self.time_actor.GetTextProperty().BoldOn()
 
     def _configure_scalar_bar(self):
-        if self._colorbar_added:
-            scalar_bar = self.plotter.scalar_bar
-            scalar_bar.SetOrientationToVertical()
-            scalar_bar.SetHeight(0.6)
-            scalar_bar.SetWidth(0.05)
-            scalar_bar.SetPosition(0.02, 0.2)
+        if self._scalar_bar is not None:
+            self._scalar_bar.SetOrientationToVertical()
+            self._scalar_bar.SetHeight(0.6)
+            self._scalar_bar.SetWidth(0.05)
+            self._scalar_bar.SetPosition(0.02, 0.2)
 
     def _configure_dock_time_widget(self, layout=None):
         len_time = len(self._data['time']) - 1
@@ -1947,12 +1942,11 @@ def add_data(self, array, fmin=None, fmid=None, fmax=None,
                 )
                 self._data['time_actor'] = time_actor
                 self._time_label_added = True
-            if colorbar and not self._colorbar_added and do:
+            if colorbar and self._scalar_bar is None and do:
                 kwargs = dict(source=actor, n_labels=8, color=self._fg_color,
                               bgcolor=self._brain_color[:3])
                 kwargs.update(colorbar_kwargs or {})
-                self._renderer.scalarbar(**kwargs)
-                self._colorbar_added = True
+                self._scalar_bar = self._renderer.scalarbar(**kwargs)
             self._renderer.set_camera(**views_dicts[hemi][v])
 
         # 4) update the scalar bar and opacity
@@ -2650,9 +2644,6 @@ def update_lut(self, fmin=None, fmid=None, fmax=None, alpha=None):
         # update our values
         rng = self._cmap_range
         ctable = self._data['ctable']
-        # in testing, no plotter; if colorbar=False, no scalar_bar
-        scalar_bar = getattr(
-            getattr(self._renderer, 'plotter', None), 'scalar_bar', None)
         for hemi in ['lh', 'rh', 'vol']:
             hemi_data = self._data.get(hemi)
             if hemi_data is not None:
@@ -2663,25 +2654,22 @@ def update_lut(self, fmin=None, fmid=None, fmax=None, alpha=None):
                                         opacity=alpha,
                                         rng=rng)
                     self._renderer._set_colormap_range(
-                        mesh._actor, ctable, scalar_bar, rng,
+                        mesh._actor, ctable, self._scalar_bar, rng,
                         self._brain_color)
-                    scalar_bar = None
 
                 grid_volume_pos = hemi_data.get('grid_volume_pos')
                 grid_volume_neg = hemi_data.get('grid_volume_neg')
                 for grid_volume in (grid_volume_pos, grid_volume_neg):
                     if grid_volume is not None:
                         self._renderer._set_volume_range(
                             grid_volume, ctable, hemi_data['alpha'],
-                            scalar_bar, rng)
-                        scalar_bar = None
+                            self._scalar_bar, rng)
 
                 glyph_actor = hemi_data.get('glyph_actor')
                 if glyph_actor is not None:
                     for glyph_actor_ in glyph_actor:
                         self._renderer._set_colormap_range(
-                            glyph_actor_, ctable, scalar_bar, rng)
-                        scalar_bar = None
+                            glyph_actor_, ctable, self._scalar_bar, rng)
         if self.time_viewer:
             with self._no_lut_update(f'update_lut {args}'):
                 for key in ('fmin', 'fmid', 'fmax'):

mne/viz/_brain/tests/test_brain.py

@@ -618,7 +618,7 @@ def test_brain_traces(renderer_interactive_pyvista, hemi, src, tmpdir,
     assert brain.traces_mode == 'vertex'
     assert hasattr(brain, "picked_points")
     assert hasattr(brain, "_spheres")
-    assert brain.plotter.scalar_bar.GetNumberOfLabels() == 3
+    assert brain._scalar_bar.GetNumberOfLabels() == 3
 
     # add foci should work for volumes
     brain.add_foci([[0, 0, 0]], hemi='lh' if src == 'surface' else 'vol')

mne/viz/backends/_pyvista.py

@@ -607,15 +607,21 @@ def text3d(self, x, y, z, text, scale, color='white'):
 
     def scalarbar(self, source, color="white", title=None, n_labels=4,
                   bgcolor=None, **extra_kwargs):
+        if isinstance(source, vtk.vtkMapper):
+            mapper = source
+        elif isinstance(source, vtk.vtkActor):
+            mapper = source.GetMapper()
+        else:
+            mapper = None
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=FutureWarning)
             kwargs = dict(color=color, title=title, n_labels=n_labels,
                           use_opacity=False, n_colors=256, position_x=0.15,
                           position_y=0.05, width=0.7, shadow=False, bold=True,
                           label_font_size=22, font_family=self.font_family,
-                          background_color=bgcolor)
+                          background_color=bgcolor, mapper=mapper)
             kwargs.update(extra_kwargs)
-            self.plotter.add_scalar_bar(**kwargs)
+            return self.plotter.add_scalar_bar(**kwargs)
 
     def show(self):
         self.plotter.show()