from __future__ import annotations
import itertools
from typing import Any
from warnings import warn
import weakref
import numpy as np
import pygfx
from imgui_bundle import imgui
from ..graphics import Graphic
from ..graphics.features import BufferManager, VertexColors, parse_dash_pattern
from ..graphics.features._scatter import user_input_to_marker
from ._base import ImguiContainer, ImguiWindow
from ._colorbar import ImguiColorbar
# marker shapes a legend can draw, the rest of the pygfx marker shapes are not made of imgui draw
# list primitives
MARKER_SHAPES = (
"circle",
"ring",
"square",
"diamond",
"plus",
"cross",
"asterisk6",
"asterisk8",
"tick",
"tick_left",
"tick_right",
"triangle_up",
"triangle_down",
"triangle_left",
"triangle_right",
)
# angles of the legs of the asterisks, in radians
ASTERISK6_ANGLES = np.arange(3) * np.pi / 3
ASTERISK8_ANGLES = np.arange(4) * np.pi / 4
SWATCH_WIDTH = 28 # width in pixels of the swatch drawn to the left of a label
SWATCH_GAP = 6 # gap in pixels between the swatch and its label
INDENT = 8 # indent in pixels of the elements drawn below an item's label
def check_marker(marker: str) -> str:
"""the pygfx marker shape name for ``marker``, raises if a legend cannot draw it"""
marker = user_input_to_marker(marker)
if marker not in MARKER_SHAPES:
raise ValueError(
f"a legend cannot draw the marker: {marker!r}, it is not made of imgui draw list "
f"primitives. The markers a legend can draw are: {MARKER_SHAPES}"
)
return marker
def imgui_color(color: pygfx.Color) -> int:
"""a color as an imgui packed color"""
return imgui.color_convert_float4_to_u32(tuple(pygfx.Color(color)))
def bar_corners(
cx: float, cy: float, length: float, width: float, angle: float
) -> tuple:
"""the corners of a ``length`` by ``width`` bar centered on (cx, cy), rotated by ``angle``"""
along = 0.5 * length * np.array([np.cos(angle), np.sin(angle)])
across = 0.5 * width * np.array([-np.sin(angle), np.cos(angle)])
center = np.array([cx, cy])
corners = [
center - along - across,
center + along - across,
center + along + across,
center - along + across,
]
return tuple((float(x), float(y)) for x, y in corners)
def draw_marker(
draw_list,
marker: str,
cx: float,
cy: float,
size: float,
color,
edge_color,
edge_width: float,
):
"""
Draw a scatter marker centered on (cx, cy), with the geometry of the pygfx marker SDFs.
``marker`` must be one of ``MARKER_SHAPES``.
"""
fill = imgui_color(color)
edge = imgui_color(edge_color)
radius = 0.5 * size
match marker:
case "circle":
draw_list.add_circle_filled((cx, cy), radius, fill)
if edge_width > 0:
draw_list.add_circle((cx, cy), radius, edge, thickness=edge_width)
case "ring":
# the difference of a disc of size/2 and a disc of size/4
draw_list.add_circle((cx, cy), 0.375 * size, fill, thickness=0.25 * size)
case "square":
draw_list.add_rect_filled(
(cx - radius, cy - radius), (cx + radius, cy + radius), fill
)
if edge_width > 0:
draw_list.add_rect(
(cx - radius, cy - radius),
(cx + radius, cy + radius),
edge,
thickness=edge_width,
)
case "diamond":
draw_list.add_quad_filled(
(cx, cy - radius),
(cx + radius, cy),
(cx, cy + radius),
(cx - radius, cy),
fill,
)
case "plus" | "cross":
# two bars of size by size/3, a cross is a plus rotated by 45 degrees
angle = 0.0 if marker == "plus" else 0.25 * np.pi
for a in (angle, angle + 0.5 * np.pi):
draw_list.add_quad_filled(*bar_corners(cx, cy, size, size / 3, a), fill)
case "asterisk6" | "asterisk8":
# legs of size by size/5
angles = ASTERISK6_ANGLES if marker == "asterisk6" else ASTERISK8_ANGLES
for a in angles:
draw_list.add_quad_filled(*bar_corners(cx, cy, size, size / 5, a), fill)
case "tick" | "tick_left" | "tick_right":
# a tick has no fill, only its edge is visible, and it is a half tick on one side
top = cy - radius if marker != "tick_right" else cy
bottom = cy + radius if marker != "tick_left" else cy
draw_list.add_line((cx, top), (cx, bottom), edge, max(edge_width, 1.0))
case "triangle_up" | "triangle_down" | "triangle_left" | "triangle_right":
# the base spans `size` and the apex is size/2 from it, centered on size/4
half = 0.5 * radius
match marker:
case "triangle_up":
points = [
(cx - radius, cy + half),
(cx + radius, cy + half),
(cx, cy - half),
]
case "triangle_down":
points = [
(cx - radius, cy - half),
(cx + radius, cy - half),
(cx, cy + half),
]
case "triangle_left":
points = [
(cx + half, cy - radius),
(cx + half, cy + radius),
(cx - half, cy),
]
case _:
points = [
(cx - half, cy - radius),
(cx - half, cy + radius),
(cx + half, cy),
]
draw_list.add_triangle_filled(*points, fill)
if edge_width > 0:
draw_list.add_triangle(*points, edge, thickness=edge_width)
[docs]
class LegendElement:
"""One entry of a :class:`.LegendItem`, a swatch and its label."""
def __init__(self, label: str, color):
self._label = label
self._color = pygfx.Color(color)
@property
def label(self) -> str:
"""the label drawn next to the swatch"""
return self._label
@property
def color(self) -> pygfx.Color:
"""the color the swatch is drawn with"""
return self._color
@property
def swatch_height(self) -> float:
"""height in pixels the swatch needs, must be implemented in subclass"""
raise NotImplementedError
[docs]
def draw_swatch(self, draw_list, x: float, y: float, width: float, height: float):
"""draw the swatch in the given rect, must be implemented in subclass"""
raise NotImplementedError
def __repr__(self) -> str:
return f"{type(self).__name__} '{self.label}'"
[docs]
class LineLegendElement(LegendElement):
def __init__(self, label: str, color, thickness: float, dash_pattern):
"""A line sample drawn with the thickness and dash pattern of the graphic."""
super().__init__(label, color)
self._thickness = float(thickness)
self._dash_pattern = parse_dash_pattern(dash_pattern)
@property
def swatch_height(self) -> float:
return self._thickness
[docs]
def draw_swatch(self, draw_list, x: float, y: float, width: float, height: float):
color = imgui_color(self._color)
cy = y + 0.5 * height
# the dash pattern is the length of alternating strokes and gaps, in units of the
# thickness. An odd pattern repeats so that strokes and gaps alternate.
pattern = np.asarray(self._dash_pattern, dtype=float) * self._thickness
if pattern.size % 2:
pattern = np.tile(pattern, 2)
if pattern.sum() <= 0:
# no pattern, or one that never advances along the line
draw_list.add_line((x, cy), (x + width, cy), color, self._thickness)
return
position = x
for stroke, gap in itertools.cycle(pattern.reshape(-1, 2)):
if position >= x + width:
break
if stroke == 0:
# a zero length stroke is a dot
draw_list.add_circle_filled(
(position, cy), 0.5 * self._thickness, color
)
else:
end = min(position + stroke, x + width)
draw_list.add_line((position, cy), (end, cy), color, self._thickness)
position += stroke + gap
[docs]
class ScatterLegendElement(LegendElement):
def __init__(
self,
label: str,
color,
marker: str,
size: float,
edge_color,
edge_width: float,
):
"""A marker drawn with the marker shape, size and edge of the graphic."""
super().__init__(label, color)
self._marker = check_marker(marker)
self._size = float(size)
self._edge_color = pygfx.Color(edge_color)
self._edge_width = float(edge_width)
@property
def swatch_height(self) -> float:
return self._size
[docs]
def draw_swatch(self, draw_list, x: float, y: float, width: float, height: float):
draw_marker(
draw_list,
self._marker,
x + 0.5 * width,
y + 0.5 * height,
self._size,
self._color,
self._edge_color,
self._edge_width,
)
[docs]
class LegendItem(ImguiContainer):
"""
The legend representation of a graphic, created with ``Graphic.create_legend_item()``.
An item has one :class:`.LegendElement` per labelled value of the graphic, or an
:class:`.ImguiColorbar` when the graphic uses a quantitative colormap. It follows the graphic:
when a feature it represents changes, its elements are made again.
"""
COLORBAR_HEIGHT = 150 # height in pixels of the colorbar of a quantitative colormap
# the graphic features the elements are made from, the elements are made again when one of
# them changes
_events: tuple[str, ...] = ()
def __init__(
self, graphic: Graphic, label: str = None, cmap_transform_labels: dict = None
):
super().__init__()
# the graphic owns its legend item, hold it weakly so they are not in a reference cycle
self._graphic = weakref.ref(graphic)
self._label = label
self._cmap_transform_labels = cmap_transform_labels
self._legend: Legend | None = None
self._colorbar: ImguiColorbar | None = None
self._elements: list[LegendElement] = list()
# switching between uniform colors, per-vertex colors and a colormap replaces the feature
# instance, which cannot be followed with an event, see update()
self._colors_feature = graphic._colors
self._cmap_feature = graphic._cmap
self._make_item()
self._connect()
@property
def graphic(self) -> Graphic | None:
"""the graphic this item represents"""
return self._graphic()
@property
def legend(self) -> Legend | None:
"""the legend this item has been added to"""
return self._legend
@property
def elements(self) -> tuple[LegendElement, ...]:
"""the elements of this item, empty when it is shown as a colorbar"""
return tuple(self._elements)
@property
def colorbar(self) -> ImguiColorbar | None:
"""the colorbar of this item, set when the graphic uses a quantitative colormap"""
return self._colorbar
@property
def label(self) -> str | None:
"""get or set the label of this item, the graphic's name is used if it is not set"""
if self._label is not None:
return self._label
return self.graphic.name
@label.setter
def label(self, value: str | None):
self._label = value
self._make_item()
@property
def cmap_transform_labels(self) -> dict | None:
"""get or set the label of each cmap_transform value of a qualitative colormap"""
return self._cmap_transform_labels
@cmap_transform_labels.setter
def cmap_transform_labels(self, labels: dict | None):
self._cmap_transform_labels = labels
self._make_item()
def _fpl_add_hook(self, figure):
super()._fpl_add_hook(figure)
if self._colorbar is not None:
self._colorbar._fpl_add_hook(figure)
def _make_item(self):
"""make the colorbar or the elements that represent the graphic's current state"""
self._colorbar = None
self._elements = list()
cmap = self.graphic.cmap
if cmap is not None and cmap.interpolation != "nearest":
# a quantitative colormap is shown as a colorbar rather than as elements
self._colorbar = ImguiColorbar(
graphics=self.graphic, title=self.label, height=self.COLORBAR_HEIGHT
)
if self._figure is not None:
self._colorbar._fpl_add_hook(self._figure)
return
self._elements = self._make_elements()
def _remake_item(self):
"""make the item again after the graphic changed"""
try:
self._make_item()
except ValueError as exc:
# the graphic changed into something the labels it has cannot represent, the user can
# set the labels it needs now
warn(f"the legend item of {self.graphic} cannot be made: {exc}")
def _make_elements(self) -> list[LegendElement]:
"""the elements representing the graphic, must be implemented in subclass"""
raise NotImplementedError
def _get_label(self) -> str:
"""the label of the item's single element"""
label = self.label
if label is None:
raise ValueError(
"pass a `label` to `create_legend_item()`, or give the graphic a `name` to use as "
"its legend label"
)
return label
def _get_color(self) -> pygfx.Color:
"""the color an element is drawn with when it is not a colormap color itself"""
graphic = self.graphic
if graphic.cmap is None:
return graphic.colors
# a colormap has no single color, use the color of the first datapoint
return self._get_cmap_color(graphic.cmap_transform[0])
def _get_cmap_color(self, value: float) -> pygfx.Color:
"""the color a cmap_transform value maps onto, i.e. what the material's maprange gives it"""
graphic = self.graphic
vmin, vmax = graphic.cmap_range
span = vmax - vmin
return pygfx.Color(
np.asarray(graphic.cmap((value - vmin) / span if span else 0.0))
)
def _get_cmap_colors(self) -> list[tuple[pygfx.Color, str]]:
"""(color, label) of each labelled cmap_transform value, empty without a colormap"""
if self.graphic.cmap is None:
return list()
if self._cmap_transform_labels is None:
raise ValueError(
"a qualitative colormap needs `cmap_transform_labels`, a dict mapping each "
"cmap_transform value onto its legend label"
)
return [
(self._get_cmap_color(value), label)
for value, label in self._cmap_transform_labels.items()
]
@staticmethod
def _get_labelled_values(
feature, value, labels: dict, name: str
) -> list[tuple[Any, str]]:
"""
(value, label) of each labelled value of a feature.
A per-datapoint feature needs a label for each of its values. A uniform feature has the
single ``value`` and is labelled only when ``labels`` is given.
"""
if isinstance(feature, BufferManager):
if labels is None:
raise ValueError(
f"`{name}` is per-datapoint, pass `{name}_labels`, a dict mapping each "
f"{name} value onto its legend label"
)
return list(labels.items())
if labels is None:
return list()
if value not in labels:
raise ValueError(
f"`{name}_labels` has no label for the graphic's {name}: {value!r}"
)
return [(value, labels[value])]
@staticmethod
def _get_feature_value(feature):
"""a feature's value, the first datapoint's value if it is per-datapoint"""
if isinstance(feature, BufferManager):
return feature.value[0]
return feature.value
def _check_vertex_colors(self) -> bool:
"""whether the graphic uses per-vertex colors, which a legend cannot represent"""
if not isinstance(self.graphic._colors, VertexColors):
return False
warn(
f"{self.graphic} uses per-vertex colors, which a legend item cannot represent. Use a "
f"uniform color or a colormap."
)
return True
def _connect(self):
"""subscribe to the events of the features the elements are made from"""
graphic = self.graphic
for event in (*self._events, "name", "deleted"):
if getattr(graphic, f"_{event}", None) is None:
# the feature does not exist in the graphic's current mode, e.g. colors while a
# colormap is set
continue
graphic.add_event_handler(self._graphic_event_handler, event)
def _disconnect(self):
"""disconnect the event handlers of the graphic"""
graphic = self.graphic
for event, handlers in graphic.event_handlers:
if self._graphic_event_handler in handlers:
graphic.remove_event_handler(self._graphic_event_handler, event)
def _graphic_event_handler(self, ev):
"""the graphic changed, make the item again"""
if ev.type == "deleted":
if self._legend is not None:
self._legend.remove(self)
return
self._remake_item()
@property
def _draws_label(self) -> bool:
"""
whether this item draws its label above its elements, which takes whole lines
A single element is labelled by itself, several of them are labelled by the item.
"""
return len(self._elements) > 1 and self.label is not None
[docs]
def check_graphic_mode(self):
"""
Make the item again if the graphic switched between uniform colors, per-vertex colors and
a colormap, which replaces the feature instance and clears its handlers, so there is no
event to follow. Called on each draw.
"""
graphic = self.graphic
if (graphic._colors is self._colors_feature) and (
graphic._cmap is self._cmap_feature
):
return
self._colors_feature = graphic._colors
self._cmap_feature = graphic._cmap
self._disconnect()
self._connect()
self._remake_item()
[docs]
def update(self):
draw_items([self])
def __repr__(self) -> str:
return f"{type(self).__name__} of {self.graphic}"
[docs]
class LineLegendItem(LegendItem):
_events = (
"colors",
"cmap",
"cmap_transform",
"cmap_range",
"thickness",
"dash_pattern",
)
def __init__(
self,
graphic,
label: str = None,
dash_pattern_labels: dict = None,
cmap_transform_labels: dict = None,
):
"""
The legend representation of a :class:`.LineGraphic`, created with
``LineGraphic.create_legend_item()``.
"""
self._dash_pattern_labels = self._parse_dash_pattern_labels(dash_pattern_labels)
super().__init__(
graphic, label=label, cmap_transform_labels=cmap_transform_labels
)
@staticmethod
def _parse_dash_pattern_labels(labels: dict | None) -> dict | None:
"""the labels keyed by the parsed dash pattern, so "--" and (5, 5) are the same key"""
if labels is None:
return None
return {parse_dash_pattern(pattern): label for pattern, label in labels.items()}
@property
def dash_pattern_labels(self) -> dict | None:
"""get or set the label of each dash pattern"""
return self._dash_pattern_labels
@dash_pattern_labels.setter
def dash_pattern_labels(self, labels: dict | None):
self._dash_pattern_labels = self._parse_dash_pattern_labels(labels)
self._make_item()
def _make_elements(self) -> list[LegendElement]:
graphic = self.graphic
if self._check_vertex_colors():
return list()
color = self._get_color()
thickness = graphic.thickness
dash_pattern = parse_dash_pattern(graphic.dash_pattern)
elements = [
LineLegendElement(label, cmap_color, thickness, dash_pattern)
for cmap_color, label in self._get_cmap_colors()
]
elements += [
LineLegendElement(label, color, thickness, value)
for value, label in self._get_labelled_values(
graphic._dash_pattern,
dash_pattern,
self._dash_pattern_labels,
"dash_pattern",
)
]
if elements:
return elements
return [LineLegendElement(self._get_label(), color, thickness, dash_pattern)]
[docs]
class ScatterLegendItem(LegendItem):
_events = (
"colors",
"cmap",
"cmap_transform",
"cmap_range",
"markers",
"sizes",
"edge_colors",
"edge_width",
)
def __init__(
self,
graphic,
label: str = None,
markers_labels: dict = None,
sizes_labels: dict = None,
cmap_transform_labels: dict = None,
):
"""
The legend representation of a :class:`.ScatterGraphic`, created with
``ScatterGraphic.create_legend_item()``.
"""
if graphic.mode != "markers":
raise ValueError(
f"a legend item is only supported for a scatter with mode 'markers', the mode of "
f"this scatter is: {graphic.mode!r}"
)
self._markers_labels = self._parse_markers_labels(markers_labels)
self._sizes_labels = sizes_labels
super().__init__(
graphic, label=label, cmap_transform_labels=cmap_transform_labels
)
@staticmethod
def _parse_markers_labels(labels: dict | None) -> dict | None:
"""the labels keyed by the pygfx marker shape name, so "o" and "circle" are the same key"""
if labels is None:
return None
return {check_marker(marker): label for marker, label in labels.items()}
@property
def markers_labels(self) -> dict | None:
"""get or set the label of each marker"""
return self._markers_labels
@markers_labels.setter
def markers_labels(self, labels: dict | None):
self._markers_labels = self._parse_markers_labels(labels)
self._make_item()
@property
def sizes_labels(self) -> dict | None:
"""get or set the label of each point size"""
return self._sizes_labels
@sizes_labels.setter
def sizes_labels(self, labels: dict | None):
self._sizes_labels = labels
self._make_item()
def _make_elements(self) -> list[LegendElement]:
graphic = self.graphic
if self._check_vertex_colors():
return list()
color = self._get_color()
marker = check_marker(self._get_feature_value(graphic._markers))
size = self._get_feature_value(graphic._sizes)
edge_color = self._get_feature_value(graphic._edge_colors)
edge_width = graphic.edge_width
elements = [
ScatterLegendElement(
label, cmap_color, marker, size, edge_color, edge_width
)
for cmap_color, label in self._get_cmap_colors()
]
elements += [
ScatterLegendElement(label, color, value, size, edge_color, edge_width)
for value, label in self._get_labelled_values(
graphic._markers, marker, self._markers_labels, "markers"
)
]
elements += [
ScatterLegendElement(label, color, marker, value, edge_color, edge_width)
for value, label in self._get_labelled_values(
graphic._sizes, size, self._sizes_labels, "sizes"
)
]
if elements:
return elements
return [
ScatterLegendElement(
self._get_label(), color, marker, size, edge_color, edge_width
)
]
def get_element_width(element: LegendElement) -> float:
"""width in pixels an element needs for its swatch and its label"""
return SWATCH_WIDTH + SWATCH_GAP + imgui.calc_text_size(element.label).x
def draw_element(element: LegendElement):
"""draw an element, its swatch and its label next to it"""
draw_list = imgui.get_window_draw_list()
text_height = imgui.get_text_line_height()
height = max(text_height, element.swatch_height)
pos = imgui.get_cursor_screen_pos()
element.draw_swatch(draw_list, pos.x, pos.y, SWATCH_WIDTH, height)
imgui.dummy((SWATCH_WIDTH, height))
imgui.same_line(spacing=SWATCH_GAP)
# center the label on the swatch
imgui.set_cursor_pos_y(imgui.get_cursor_pos_y() + 0.5 * (height - text_height))
imgui.text(element.label)
def draw_elements(elements: list[LegendElement], width: float):
"""
Draw elements as a grid of columns of ``width``, as many of them as fit the width available
where the grid starts. A wide window gets more columns and fewer rows, a narrow one gets one
column.
"""
if len(elements) == 0:
return
spacing = imgui.get_style().item_spacing.x
available = imgui.get_content_region_avail().x
n_columns = int((available + spacing) // (width + spacing))
n_columns = min(max(n_columns, 1), len(elements))
for i, element in enumerate(elements):
column = i % n_columns
if column:
# every column is the same width, so they line up
imgui.same_line(offset_from_start_x=column * (width + spacing))
draw_element(element)
def draw_items(items: list):
"""
Draw legend items as one grid of elements that fits the width available.
Every element gets a cell of the same width, the widest of them all, so that the columns line
up. The elements of the items that do not draw a label flow together into the grid, an item
that draws its label, and a colorbar, take whole lines.
"""
elements = [
element
for item in items
if isinstance(item, LegendItem)
for element in item.elements
]
width = max((get_element_width(element) for element in elements), default=0.0)
flowing: list[LegendElement] = list()
for item in items:
if isinstance(item, LegendItem):
if item.graphic is None:
# the graphic has been garbage collected
continue
item.check_graphic_mode()
if item.colorbar is None and not item._draws_label:
flowing += item.elements
continue
# anything that takes whole lines ends the run of flowing elements
draw_elements(flowing, width)
flowing = list()
if isinstance(item, ImguiColorbar):
item.draw()
elif item.colorbar is not None:
item.colorbar.draw()
else:
imgui.text(item.label)
imgui.indent(INDENT)
draw_elements(item.elements, width)
imgui.unindent(INDENT)
draw_elements(flowing, width)
[docs]
class Legend(ImguiWindow):
def __init__(self, items: list = None):
"""
An imgui legend.
Add the legend items of graphics and colorbars to it, imgui lays them out in the order they
were added. A legend is an imgui window, place it with ``Figure.add_imgui_window()`` or
``Subplot.add_imgui_window()``.
Parameters
----------
items: list of LegendItem | ImguiColorbar, optional
items to add to the legend
"""
super().__init__()
self._items: list[LegendItem | ImguiColorbar] = list()
if items is not None:
for item in items:
self.add(item)
@property
def items(self) -> tuple:
"""the items in this legend, in the order they are drawn"""
return tuple(self._items)
[docs]
def add(self, item: LegendItem | ImguiColorbar):
"""
Add an item to the legend.
Parameters
----------
item: LegendItem | ImguiColorbar
a graphic's legend item, created with ``Graphic.create_legend_item()``, or a colorbar
"""
if not isinstance(item, (LegendItem, ImguiColorbar)):
raise TypeError(
f"a legend takes `LegendItem` and `ImguiColorbar` instances, you have passed a: "
f"{type(item).__name__}"
)
if isinstance(item, LegendItem) and item.legend is not None:
raise ValueError(
"this legend item is already in a legend, a `LegendItem` can only be in one "
"`Legend`"
)
self._items.append(item)
if isinstance(item, LegendItem):
item._legend = self
if self._figure is not None:
item._fpl_add_hook(self._figure)
[docs]
def remove(self, item: LegendItem | ImguiColorbar):
"""
Remove an item from the legend, it can be added to a legend again later.
Parameters
----------
item: LegendItem | ImguiColorbar
the item to remove
"""
self._items.remove(item)
if isinstance(item, LegendItem):
item._legend = None
def _fpl_add_hook(self, figure, **kwargs):
# the items are drawn within this window, they need the figure too
super()._fpl_add_hook(figure, **kwargs)
for item in self._items:
item._fpl_add_hook(figure)
[docs]
def update(self):
draw_items(self._items)
def __repr__(self) -> str:
return f"Legend of <{len(self._items)}> items"
