"""PyVista Scalar bar module.""" from __future__ import annotations import contextlib import weakref import numpy as np import pyvista as pv from pyvista import MAX_N_COLOR_BARS from pyvista import _vtk from pyvista._deprecate_positional_args import _deprecate_positional_args from pyvista.core.utilities.misc import _NoNewAttrMixin from .colors import Color from .tools import parse_font_family class ScalarBars(_NoNewAttrMixin): """Plotter Scalar Bars. Parameters ---------- plotter : pyvista.Plotter Plotter that the scalar bars are associated with. """ def __init__(self, plotter): """Initialize ScalarBars.""" self._plotter = weakref.proxy(plotter) self._scalar_bar_ranges = {} self._scalar_bar_mappers = {} self._scalar_bar_actors = {} self._scalar_bar_widgets = {} def clear(self): """Remove all scalar bars and resets all scalar bar properties.""" self._scalar_bar_ranges = {} self._scalar_bar_mappers = {} self._scalar_bar_actors = {} self._scalar_bar_widgets = {} def __plotter_close__(self) -> None: """Release scalar-bar state when the owning plotter closes.""" self.clear() def __repr__(self): """Nice representation of this class.""" lines = [] lines.append('Scalar Bar Title Interactive') for title in self._scalar_bar_actors: interactive = title in self._scalar_bar_widgets title_quotes = f'"{title}"' lines.append(f'{title_quotes:20} {interactive!s:5}') return '\n'.join(lines) @_deprecate_positional_args(allowed=['actor']) def _remove_mapper_from_plotter( self, actor, reset_camera: bool = False, # noqa: FBT001, FBT002 render: bool = False, # noqa: FBT001, FBT002 ): # numpydoc ignore=PR01,RT01 """Remove an actor's mapper from the given plotter's _scalar_bar_mappers. This ensures that when actors are removed, their corresponding scalar bars are removed. """ try: mapper = actor.GetMapper() except AttributeError: return # NOTE: keys to list to prevent iterator changing during loop for name in list(self._scalar_bar_mappers): with contextlib.suppress(ValueError): self._scalar_bar_mappers[name].remove(mapper) if not self._scalar_bar_mappers[name]: slot = self._plotter._scalar_bar_slot_lookup.pop(name, None) if slot is not None: self._scalar_bar_mappers.pop(name) self._scalar_bar_ranges.pop(name) self._plotter.remove_actor( self._scalar_bar_actors.pop(name), reset_camera=reset_camera, render=render, ) self._plotter._scalar_bar_slots.add(slot) return @_deprecate_positional_args(allowed=['title']) def remove_scalar_bar(self, title=None, render: bool = True): # noqa: FBT001, FBT002 """Remove a scalar bar. Parameters ---------- title : str, optional Title of the scalar bar to remove. Required if there is more than one scalar bar. render : bool, default: True Render upon scalar bar removal. Set this to ``False`` to stop the render window from rendering when a scalar bar is removed. Examples -------- Remove a scalar bar from a plotter. >>> import pyvista as pv >>> mesh = pv.Sphere() >>> mesh['data'] = mesh.points[:, 2] >>> pl = pv.Plotter() >>> _ = pl.add_mesh(mesh, cmap='coolwarm') >>> pl.remove_scalar_bar() >>> pl.show() """ if title is None: if len(self) > 1: titles = ', '.join(f'"{key}"' for key in self._scalar_bar_actors) msg = ( 'Multiple scalar bars found. Pick title of the' f'scalar bar from one of the following:\n{titles}' ) raise ValueError(msg) else: title = next(iter(self._scalar_bar_actors.keys())) actor = self._scalar_bar_actors.pop(title) self._plotter.remove_actor(actor, render=render) self._scalar_bar_ranges.pop(title) self._scalar_bar_mappers.pop(title) # add back in the scalar bar slot slot = self._plotter._scalar_bar_slot_lookup.pop(title, None) if slot is not None: self._plotter._scalar_bar_slots.add(slot) widget = self._scalar_bar_widgets.pop(title, None) if widget is not None: widget.SetEnabled(0) def __len__(self): """Return the number of scalar bar actors.""" return len(self._scalar_bar_actors) def __getitem__(self, index): """Return a scalar bar actor.""" return self._scalar_bar_actors[index] def keys(self): # numpydoc ignore=RT01 """Scalar bar keys.""" return self._scalar_bar_actors.keys() def values(self): # numpydoc ignore=RT01 """Scalar bar values.""" return self._scalar_bar_actors.values() def items(self): # numpydoc ignore=RT01 """Scalar bar items.""" return self._scalar_bar_actors.items() def __contains__(self, key) -> bool: """Check if a title is a valid actors.""" return key in self._scalar_bar_actors def update_title( self, old_title: str, new_title: str, *, render: bool = False, ) -> None: """Update the title of an existing scalar bar. .. versionadded:: 0.48.0 Parameters ---------- old_title : str Current title of the scalar bar to update. new_title : str New title for the scalar bar. render : bool, default: False Force a render after updating the title. Raises ------ KeyError If no scalar bar with ``old_title`` exists. ValueError If a scalar bar with ``new_title`` already exists. Examples -------- Update the title of a scalar bar. >>> import pyvista as pv >>> mesh = pv.Sphere() >>> mesh['Data'] = mesh.points[:, 2] >>> pl = pv.Plotter() >>> _ = pl.add_mesh(mesh, scalars='Data') >>> pl.scalar_bars.update_title('Data', 'Elevation') >>> pl.show() """ if old_title not in self._scalar_bar_actors: msg = f'Scalar bar with title "{old_title}" not found.' raise KeyError(msg) if old_title != new_title and new_title in self._scalar_bar_actors: msg = f'Scalar bar with title "{new_title}" already exists.' raise ValueError(msg) if old_title != new_title: self._scalar_bar_actors[new_title] = self._scalar_bar_actors.pop(old_title) self._scalar_bar_ranges[new_title] = self._scalar_bar_ranges.pop(old_title) self._scalar_bar_mappers[new_title] = self._scalar_bar_mappers.pop(old_title) if old_title in self._scalar_bar_widgets: self._scalar_bar_widgets[new_title] = self._scalar_bar_widgets.pop(old_title) slot = self._plotter._scalar_bar_slot_lookup.pop(old_title, None) if slot is not None: self._plotter._scalar_bar_slot_lookup[new_title] = slot self._scalar_bar_actors[new_title].SetTitle(new_title) if render: self._plotter.render() @_deprecate_positional_args(allowed=['title']) def add_scalar_bar( # noqa: PLR0917 self, title='', mapper=None, n_labels=5, italic: bool = False, # noqa: FBT001, FBT002 bold: bool = False, # noqa: FBT001, FBT002 title_font_size=None, label_font_size=None, color=None, font_family=None, shadow: bool = False, # noqa: FBT001, FBT002 width=None, height=None, position_x=None, position_y=None, vertical=None, interactive=None, fmt=None, use_opacity: bool = True, # noqa: FBT001, FBT002 outline: bool = False, # noqa: FBT001, FBT002 nan_annotation: bool = False, # noqa: FBT001, FBT002 below_label=None, above_label=None, background_color=None, n_colors=None, fill: bool = False, # noqa: FBT001, FBT002 render: bool = False, # noqa: FBT001, FBT002 theme=None, unconstrained_font_size: bool = False, # noqa: FBT001, FBT002 unique_bar: bool = False, # noqa: FBT001, FBT002 ): """Create scalar bar using the ranges as set by the last input mesh. Parameters ---------- title : str, default: "" Title of the scalar bar. Default is rendered as an empty title. mapper : :vtk:`vtkMapper`, optional Mapper used for the scalar bar. Defaults to the last mapper created by the plotter. n_labels : int, default: 5 Number of labels to use for the scalar bar. italic : bool, default: False Italicises title and bar labels. bold : bool, default: False Bolds title and bar labels. title_font_size : float, optional Sets the size of the title font. Defaults to ``None`` and is sized according to :attr:`pyvista.plotting.themes.Theme.font`. label_font_size : float, optional Sets the size of the title font. Defaults to ``None`` and is sized according to :attr:`pyvista.plotting.themes.Theme.font`. color : ColorLike, optional Either a string, rgb list, or hex color string. Default set by :attr:`pyvista.plotting.themes.Theme.font`. Can be in one of the following formats: * ``color='white'`` * ``color='w'`` * ``color=[1.0, 1.0, 1.0]`` * ``color='#FFFFFF'`` font_family : {'courier', 'times', 'arial'} Font family. Default is set by :attr:`pyvista.plotting.themes.Theme.font`. shadow : bool, default: False Adds a black shadow to the text. width : float, optional The percentage (0 to 1) width of the window for the colorbar. Default set by :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal` depending on the value of ``vertical``. height : float, optional The percentage (0 to 1) height of the window for the colorbar. Default set by :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal` depending on the value of ``vertical``. position_x : float, optional The percentage (0 to 1) along the windows's horizontal direction to place the bottom left corner of the colorbar. Default set by :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal` depending on the value of ``vertical``. position_y : float, optional The percentage (0 to 1) along the windows's vertical direction to place the bottom left corner of the colorbar. Default set by :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal` depending on the value of ``vertical``. vertical : bool, optional Use vertical or horizontal scalar bar. Default set by :attr:`pyvista.plotting.themes.Theme.colorbar_orientation`. interactive : bool, optional Use a widget to control the size and location of the scalar bar. Default set by :attr:`pyvista.plotting.themes.Theme.interactive`. fmt : str, optional ``printf`` format for labels. Default set by :attr:`pyvista.plotting.themes.Theme.font`. use_opacity : bool, default: True Optionally display the opacity mapping on the scalar bar. outline : bool, default: False Optionally outline the scalar bar to make opacity mappings more obvious. nan_annotation : bool, default: False Annotate the NaN color. below_label : str, optional String annotation for values below the scalars range. above_label : str, optional String annotation for values above the scalars range. background_color : ColorLike, optional The color used for the background in RGB format. n_colors : int, optional The maximum number of color displayed in the scalar bar. fill : bool, default: False Draw a filled box behind the scalar bar with the ``background_color``. render : bool, default: False Force a render when True. theme : pyvista.plotting.themes.Theme, optional Plot-specific theme. By default, calling from the ``Plotter``, will use the plotter theme. Setting to ``None`` will use the global theme. unconstrained_font_size : bool, default: False Whether the font size of title and labels is unconstrained. When it is constrained, the size of the scalar bar will constrain the font size. When it is not, the size of the font will always be respected. Using custom labels will force this to be ``True``. .. versionadded:: 0.44.0 unique_bar : bool, default: False Whether to create a scalar bar which is unique to the subplot. If ``True``, the scalar bar will be created with a unique key which is not shared with other subplots, even if the input title is the same. .. note:: Scalar bars are managed by a dictionary with the title as the key. By default, if a scalar bar with the same title already exists, the scalar bar will be shared. If ``unique_bar`` is ``True``, the scalar bar will be created with a unique key which is the title suffixed with ``_UNIQUE_ID_{active_renderer_index}``, where ``active_renderer_index`` is the index of the active renderer in the plotter. This allows for multiple scalar bars with the same title to be created across different subplots. .. versionadded:: 0.48.0 Returns ------- :vtk:`vtkScalarBarActor` Scalar bar actor. Notes ----- Setting ``title_font_size``, or ``label_font_size`` disables automatic font sizing for both the title and label. See Also -------- :ref:`scalar_bar_example` Examples -------- Add a custom interactive scalar bar that is horizontal, has an outline, and has a custom formatting. >>> import pyvista as pv >>> sphere = pv.Sphere() >>> sphere['Data'] = sphere.points[:, 2] >>> pl = pv.Plotter() >>> _ = pl.add_mesh(sphere, show_scalar_bar=False) >>> _ = pl.add_scalar_bar( ... 'Data', ... interactive=True, ... vertical=False, ... title_font_size=35, ... label_font_size=30, ... outline=True, ... fmt='%10.5f', ... ) >>> pl.show() """ if mapper is None: msg = 'Mapper cannot be ``None`` when creating a scalar bar' raise ValueError(msg) if theme is None: theme = pv.global_theme if interactive is None: interactive = theme.interactive if font_family is None: font_family = theme.font.family if label_font_size is None: label_font_size = theme.font.label_size if title_font_size is None: title_font_size = theme.font.title_size if fmt is None: fmt = theme.font.fmt if vertical is None and theme.colorbar_orientation.lower() == 'vertical': vertical = True # Automatically choose size if not specified if width is None: width = theme.colorbar_vertical.width if vertical else theme.colorbar_horizontal.width if height is None: if vertical: height = theme.colorbar_vertical.height else: height = theme.colorbar_horizontal.height display_title = title if unique_bar: title = f'{title}_UNIQUE_ID_{self._plotter.renderers.active_index}' # Check that this data hasn't already been plotted if title in list(self._scalar_bar_ranges.keys()): clim = list(self._scalar_bar_ranges[title]) newrng = mapper.scalar_range oldmappers = self._scalar_bar_mappers[title] # get max for range and reset everything clim[0] = min(newrng[0], clim[0]) clim[1] = max(newrng[1], clim[1]) for mh in oldmappers: mh.scalar_range = clim[0], clim[1] mapper.scalar_range = clim[0], clim[1] self._scalar_bar_mappers[title].append(mapper) self._scalar_bar_ranges[title] = clim self._scalar_bar_actors[title].SetLookupTable(mapper.lookup_table) # Color bar already present and ready to be used so returning return None # Automatically choose location if not specified if position_x is None or position_y is None: if not self._plotter._scalar_bar_slots: msg = f'Maximum number of color bars ({MAX_N_COLOR_BARS}) reached.' raise RuntimeError(msg) slot = min(self._plotter._scalar_bar_slots) self._plotter._scalar_bar_slots.remove(slot) self._plotter._scalar_bar_slot_lookup[title] = slot if position_x is None: if vertical: position_x = theme.colorbar_vertical.position_x position_x -= slot * (width + 0.2 * width) else: position_x = theme.colorbar_horizontal.position_x if position_y is None: if vertical: position_y = theme.colorbar_vertical.position_y else: position_y = theme.colorbar_horizontal.position_y position_y += slot * height # parse color color = Color(color, default_color=theme.font.color) # Create scalar bar scalar_bar = _vtk.vtkScalarBarActor() # self._scalar_bars.append(scalar_bar) if background_color is not None: background_color = np.array(Color(background_color).int_rgba) scalar_bar.GetBackgroundProperty().SetColor(background_color[0:3]) if fill: scalar_bar.DrawBackgroundOn() lut = pv.LookupTable() lut.DeepCopy(mapper.lookup_table) ctable = _vtk.vtk_to_numpy(lut.GetTable()) alphas = ctable[:, -1][:, np.newaxis] / 255.0 use_table = ctable.copy() use_table[:, -1] = 255.0 ctable = (use_table * alphas) + background_color * (1 - alphas) lut.SetTable(_vtk.numpy_to_vtk(ctable, array_type=_vtk.VTK_UNSIGNED_CHAR)) else: lut = mapper.lookup_table scalar_bar.SetLookupTable(lut) if n_colors is None: # ensure the number of colors in the scalarbar's lookup table is at # least the number in the mapper n_colors = mapper.lookup_table.n_values scalar_bar.SetMaximumNumberOfColors(n_colors) if n_labels < 1: scalar_bar.SetDrawTickLabels(False) else: scalar_bar.SetDrawTickLabels(True) scalar_bar.SetNumberOfLabels(n_labels) if nan_annotation: scalar_bar.DrawNanAnnotationOn() if above_label is not None: scalar_bar.DrawAboveRangeSwatchOn() scalar_bar.SetAboveRangeAnnotation(above_label) elif lut.above_range_color: scalar_bar.DrawAboveRangeSwatchOn() scalar_bar.SetAboveRangeAnnotation('above') if below_label is not None: scalar_bar.DrawBelowRangeSwatchOn() scalar_bar.SetBelowRangeAnnotation(below_label) elif lut.below_range_color: scalar_bar.DrawBelowRangeSwatchOn() scalar_bar.SetBelowRangeAnnotation('below') # edit the size of the colorbar scalar_bar.SetHeight(height) scalar_bar.SetWidth(width) scalar_bar.SetPosition(position_x, position_y) if fmt is not None: scalar_bar.SetLabelFormat(fmt) if vertical: scalar_bar.SetOrientationToVertical() else: scalar_bar.SetOrientationToHorizontal() if label_font_size is not None or title_font_size is not None: scalar_bar.SetUnconstrainedFontSize(True) scalar_bar.SetAnnotationTextScaling(False) else: scalar_bar.SetAnnotationTextScaling(True) label_text = scalar_bar.GetLabelTextProperty() anno_text = scalar_bar.GetAnnotationTextProperty() label_text.SetColor(color.float_rgb) anno_text.SetColor(color.float_rgb) label_text.SetShadow(shadow) anno_text.SetShadow(shadow) # Set font label_text.SetFontFamily(parse_font_family(font_family)) anno_text.SetFontFamily(parse_font_family(font_family)) label_text.SetItalic(italic) anno_text.SetItalic(italic) label_text.SetBold(bold) anno_text.SetBold(bold) if label_font_size: label_text.SetFontSize(label_font_size) anno_text.SetFontSize(label_font_size) # Set properties self._scalar_bar_ranges[title] = mapper.scalar_range self._scalar_bar_mappers[title] = [mapper] scalar_bar.SetTitle(display_title) title_text = scalar_bar.GetTitleTextProperty() title_text.SetJustificationToCentered() title_text.SetItalic(italic) title_text.SetBold(bold) title_text.SetShadow(shadow) if title_font_size: title_text.SetFontSize(title_font_size) # Set font title_text.SetFontFamily(parse_font_family(font_family)) # set color title_text.SetColor(color.float_rgb) self._scalar_bar_actors[title] = scalar_bar if interactive: scalar_widget = _vtk.vtkScalarBarWidget() scalar_widget.SetScalarBarActor(scalar_bar) scalar_widget.SetInteractor(self._plotter.iren.interactor) scalar_widget.SetEnabled(1) rep = scalar_widget.GetRepresentation() scalar_widget.On() if vertical is True or vertical is None: rep.SetOrientation(1) # type: ignore[attr-defined] # 0 = Horizontal, 1 = Vertical else: # y position determined empirically y = -position_y / 2 - height - scalar_bar.GetPosition()[1] rep.GetPositionCoordinate().SetValue(width, y) # type: ignore[attr-defined] rep.GetPosition2Coordinate().SetValue(height, width) # type: ignore[attr-defined] rep.SetOrientation(0) # type: ignore[attr-defined] # 0 = Horizontal, 1 = Vertical self._scalar_bar_widgets[title] = scalar_widget if use_opacity: scalar_bar.SetUseOpacity(True) if outline: scalar_bar.SetDrawFrame(True) frame_prop = scalar_bar.GetFrameProperty() frame_prop.SetColor(color.float_rgb) else: scalar_bar.SetDrawFrame(False) if unconstrained_font_size: scalar_bar.SetUnconstrainedFontSize(True) # finally, add to the actor and return the scalar bar self._plotter.add_actor(scalar_bar, reset_camera=False, pickable=False, render=render) return scalar_bar