plot-io 0.1.0__py3-none-any.whl

plot_io/colorbar.py

@@ -0,0 +1,342 @@
+"""Colorbars: the color key of a mapped series or of a shared ``ColorMap``."""
+
+from typing import Any
+
+from pdf_io import LinearGradient
+
+from .axis import Axis, AxisState
+from .numbers import number_settings
+from .scales import get_scale
+from .settings import Setting, Settings, boolean, color, number, one_of, route, table, text
+from .style import resolve_color
+from .text import box, draw_text, measure
+
+__all__ = ["Colorbar", "BarState"]
+
+
+class Colorbar(Settings):
+    """A colorbar: attached to a plot (``plot.colorbar(…)``) or in its own grid area
+    (``figure.colorbar(…, id=…)``).  Its axis is ``colorbar.axis``."""
+
+    SETTINGS = table(
+        position=Setting(one_of("right", "left", "top", "bottom"),
+                         doc="Side of the plot an attached bar takes."),
+        orientation=Setting(one_of("vertical", "horizontal"),
+                            doc="Bar direction; default: vertical, or horizontal for bars "
+                                "attached at the top or bottom."),
+        size=Setting(number(positive=True),
+                     doc="Bar thickness, in the figure's unit; default 10 pt."),
+        length=Setting(number(lo=0.05, hi=1.0),
+                       doc="Bar length as a fraction of the axes side (attached bars)."),
+        pad=Setting(number(lo=0),
+                    doc="Gap between the axes and an attached bar, in the figure's unit."),
+        label=Setting(text, doc="Bar label: a string or rich content."),
+        extend=Setting(one_of("auto", "neither", "min", "max", "both"),
+                       doc="Triangular ends in the under/over colors: `auto` shows them when "
+                           "the limits clip the data."),
+        frame=Setting(boolean, doc="Outline the bar."),
+        frame_color=Setting(color, doc="Outline color."),
+        frame_width=Setting(number(lo=0), doc="Outline width (pt).", length=True),
+        font_size=Setting(number(positive=True), inherited=True, doc="Base font size (pt).", length=True),
+        font_family=Setting(str, inherited=True, doc="Font family."),
+        text_color=Setting(color, inherited=True, doc="Text color."),
+        **number_settings(),
+    )
+
+    def __init__(self, source: Any = None, plot: Any = None, figure: Any = None,
+                 **settings: Any) -> None:
+        axis_kw = {k: settings.pop(k) for k in list(settings) if k in Axis.SETTINGS
+                   and k not in Colorbar.SETTINGS}
+        own, = route(settings, [Colorbar])
+        self.source = source
+        self.plot = plot
+        self.figure = figure if figure is not None else (plot.figure if plot is not None else None)
+        self.id: Any = None
+        self.axis = Axis(self, "c")
+        super().__init__(**own)
+        if axis_kw:
+            self.axis.set(**axis_kw)
+
+    def _parent(self) -> Settings | None:
+        return self.plot if self.plot is not None else self.figure
+
+    def _points_per_unit(self) -> float:
+        return self.figure.points_per_unit if self.figure is not None else 72.0
+
+    def __repr__(self) -> str:
+        return f"Colorbar({self.source!r})"
+
+    def color_map(self) -> Any:
+        """The ``ColorMap`` shown (``None`` when no source is found)."""
+        from .maps import ColorMap
+
+        src = self.source
+        if src is None and self.plot is not None:
+            for s in reversed(self.plot.series):
+                if s.color_sources():
+                    src = s
+                    break
+        if src is None:
+            return None
+        if isinstance(src, ColorMap):
+            return src
+        sources = src.color_sources()
+        if not sources:
+            raise ValueError(f"{src!r} has no color map to show")
+        return sources[0][0]
+
+    def auto_label(self) -> Any:
+        """The label: the setting, else the source series' color column name, else the
+        name of a density it maps."""
+        label = self.get("label")
+        if label is not None:
+            return label
+        src = self.source
+        if src is None and self.plot is not None:
+            src = next((s for s in reversed(self.plot.series) if s.color_sources()), None)
+        if src is None or not hasattr(src, "SETTINGS"):
+            return None
+        cols = getattr(src, "columns", {})
+        for name in ("color", "line_color", "marker_color"):
+            if name in cols:
+                return cols[name]
+        from .density import Density
+        for name in ("color", "line_color", "marker_color"):
+            if name in src.SETTINGS and isinstance(src.get(name), Density):
+                return src.get(name).label()
+        return None
+
+    @property
+    def vertical(self) -> bool:
+        return self.get("orientation") == "vertical"
+
+
+class BarState:
+    """A colorbar while drawing."""
+
+    def __init__(self, cb: Colorbar, maps: dict[int, tuple], plot: Any) -> None:
+        self.cb = cb
+        self.cmap = cb.color_map()
+        if self.cmap is None:
+            raise ValueError("colorbar without a source: give a mapped series or a ColorMap")
+        flo, fhi, dlo, dhi = maps[id(self.cmap)]
+        self.state = AxisState(cb.axis)
+        scale = self.state.scale = get_scale(self.cmap.scale)
+        self.state.flo, self.state.fhi = flo, fhi
+        self.state.lo, self.state.hi = scale.inverse(flo), scale.inverse(fhi)
+        self.data = (scale.f(dlo), scale.f(dhi))
+        self.position = cb.get("position")
+        self.vertical = cb.vertical
+        self.thickness = cb.get("size") * cb._points_per_unit()
+        self.pad = cb.get("pad") * cb._points_per_unit()
+        ext = cb.get("extend")
+        if ext == "auto":
+            lo_clip = self.data[0] < flo - 1e-12 * abs(fhi - flo)
+            hi_clip = self.data[1] > fhi + 1e-12 * abs(fhi - flo)
+            ext = {(True, True): "both", (True, False): "min", (False, True): "max"}.get(
+                (lo_clip, hi_clip), "neither")
+        self.extend = ext
+        self.side = self._side()
+
+    def _side(self) -> str:
+        """Side of the bar that carries ticks and labels."""
+        if self.cb.plot is not None:
+            return self.position
+        return "right" if self.vertical else "bottom"
+
+    # -- measurement ----------------------------------------------------------------
+
+    def _labels_extent(self, canvas: Any, length: float) -> tuple[float, float]:
+        """Space outside the bar for ticks, labels and the bar label, and the overhang of
+        end labels along the bar."""
+        st = self.state
+        st.compute_ticks(canvas, length, not self.vertical)
+        axis = self.cb.axis
+        out = st.tick_out()[0]
+        font = st.tick_font()
+        labels = [t.label for t in st.ticks if t.major and t.label is not None]
+        across = out
+        over = 0.0
+        if labels:
+            sizes = [box(canvas, lab, font) for lab in labels]
+            across += axis.get("tick_pad") + max((w if self.vertical else h) for w, h in sizes)
+            over = max((h if self.vertical else w) for w, h in sizes) / 2
+        label = self.cb.auto_label()
+        if label is not None:
+            _, a, d = measure(canvas, label, st.label_font())
+            across += axis.get("label_pad") + a + d
+        return across, over
+
+    def measure_attached(self, canvas: Any, axes_len: float) -> float:
+        length = axes_len * self.cb.get("length")
+        across, _ = self._labels_extent(canvas, length)
+        return self.pad + self.thickness + across
+
+    def measure(self, canvas: Any, w: float, h: float) -> tuple[float, float, float, float]:
+        """Margins of a bar in its own area."""
+        length = h if self.vertical else w
+        across, over = self._labels_extent(canvas, max(length, 1.0))
+        tri = self.thickness if self.extend != "neither" else 0.0
+        over = max(over, tri)
+        if self.vertical:
+            return (0.0, across, over, over)
+        return (over, over, 0.0, across)
+
+    def bar_rect(self, rect: tuple[float, float, float, float]) -> tuple:
+        x0, y0, x1, y1 = rect
+        if self.vertical:
+            return (x0, y0, min(x1, x0 + self.thickness), y1)
+        return (x0, y0, x1, min(y1, y0 + self.thickness))
+
+    # -- drawing --------------------------------------------------------------------
+
+    def draw_attached(self, canvas: Any, axes: tuple, offset: float, base: float) -> float:
+        x0, y0, x1, y1 = axes
+        frac = self.cb.get("length")
+        side = self.position
+        start = base + offset + self.pad
+        if side in ("left", "right"):
+            length = (y1 - y0) * frac
+            by0 = y0 + (y1 - y0 - length) / 2
+            bx0 = x1 + start if side == "right" else x0 - start - self.thickness
+            rect = (bx0, by0, bx0 + self.thickness, by0 + length)
+        else:
+            length = (x1 - x0) * frac
+            bx0 = x0 + (x1 - x0 - length) / 2
+            by0 = y1 + start if side == "bottom" else y0 - start - self.thickness
+            rect = (bx0, by0, bx0 + length, by0 + self.thickness)
+        self.draw_area(canvas, rect)
+        across, _ = self._labels_extent(canvas, length)
+        return self.pad + self.thickness + across
+
+    def draw_area(self, canvas: Any, rect: tuple[float, float, float, float]) -> None:
+        x0, y0, x1, y1 = rect
+        cmap = self.cmap
+        vertical = self.vertical
+        length = (y1 - y0) if vertical else (x1 - x0)
+        n = 256 if cmap.interpolation == "linear" else len(cmap.stops)
+        stops: list[tuple[float, Any, float]] = []
+        if cmap.interpolation == "linear":
+            for i in range(n):
+                t = i / (n - 1)
+                c = cmap(t)
+                stops.append((t, c[:3], c[3]))
+        else:
+            for i, (p, c) in enumerate(cmap.stops):
+                nxt = cmap.stops[i + 1][0] if i + 1 < len(cmap.stops) else 1.0
+                stops.append((p, c[:3], c[3]))
+                stops.append((nxt, c[:3], c[3]))
+        inv = self.cb.axis.get("invert")
+        if vertical:
+            g = LinearGradient(0, y0, 0, y1, stops) if inv else LinearGradient(0, y1, 0, y0, stops)
+        else:
+            g = LinearGradient(x1, 0, x0, 0, stops) if inv else LinearGradient(x0, 0, x1, 0, stops)
+        canvas.save()
+        canvas.fill_style = g
+        canvas.begin_path()
+        canvas.rect(x0, y0, x1 - x0, y1 - y0)
+        canvas.fill()
+        self._extensions(canvas, rect)
+        if self.cb.get("frame"):
+            canvas.stroke_style = resolve_color(self.cb.get("frame_color"), ())
+            canvas.line_width = self.cb.get("frame_width")
+            canvas.begin_path()
+            canvas.rect(x0, y0, x1 - x0, y1 - y0)
+            canvas.stroke()
+        canvas.restore()
+        self._axis(canvas, rect, length)
+
+    def _extensions(self, canvas: Any, rect: tuple) -> None:
+        if self.extend == "neither":
+            return
+        x0, y0, x1, y1 = rect
+        t = self.thickness
+        cmap = self.cmap
+        inv = self.cb.axis.get("invert")
+        ends = []
+        if self.extend in ("min", "both"):
+            ends.append(("min", cmap.under or cmap(0.0)))
+        if self.extend in ("max", "both"):
+            ends.append(("max", cmap.over or cmap(1.0)))
+        for which, c in ends:
+            at_start = (which == "min") != inv     # start: bottom (vertical) or left
+            canvas.fill_style = c
+            canvas.begin_path()
+            if self.vertical:
+                if at_start:
+                    canvas.move_to(x0, y1)
+                    canvas.line_to(x1, y1)
+                    canvas.line_to((x0 + x1) / 2, y1 + t)
+                else:
+                    canvas.move_to(x0, y0)
+                    canvas.line_to(x1, y0)
+                    canvas.line_to((x0 + x1) / 2, y0 - t)
+            else:
+                if at_start:
+                    canvas.move_to(x0, y0)
+                    canvas.line_to(x0, y1)
+                    canvas.line_to(x0 - t, (y0 + y1) / 2)
+                else:
+                    canvas.move_to(x1, y0)
+                    canvas.line_to(x1, y1)
+                    canvas.line_to(x1 + t, (y0 + y1) / 2)
+            canvas.close_path()
+            canvas.fill()
+
+    def _axis(self, canvas: Any, rect: tuple, length: float) -> None:
+        st = self.state
+        if abs(st.length - length) > 0.01:
+            st.compute_ticks(canvas, length, not self.vertical)
+        axis = self.cb.axis
+        x0, y0, x1, y1 = rect
+        side = self.side
+        mapper = st.mapper(y1, y0) if self.vertical else st.mapper(x0, x1)
+        out, inside = st.tick_out()
+        ml = axis.get("minor_tick_length")
+        canvas.save()
+        canvas.stroke_style = resolve_color(axis.get("tick_color"), ())
+        canvas.line_width = axis.get("tick_width")
+        canvas.begin_path()
+        base = {"right": x1, "left": x0, "bottom": y1, "top": y0}[side]
+        other = {"right": x0, "left": x1, "bottom": y0, "top": y1}[side]
+        sign = 1.0 if side in ("right", "bottom") else -1.0
+        for t in st.ticks:
+            p = mapper(t.value)
+            o, i = (out, inside) if t.major else ((ml if out else 0.0), (ml if inside else 0.0))
+            ends = [(base - sign * i, base + sign * o)]
+            if axis.get("mirror") and i:
+                ends.append((other + sign * i, other))
+            for a, b in ends:
+                if self.vertical:
+                    canvas.move_to(a, p)
+                    canvas.line_to(b, p)
+                else:
+                    canvas.move_to(p, a)
+                    canvas.line_to(p, b)
+        canvas.stroke()
+        canvas.restore()
+        font = st.tick_font()
+        pos = base + sign * (out + axis.get("tick_pad"))
+        widest = 0.0
+        for t in st.ticks:
+            if not t.major or t.label is None:
+                continue
+            p = mapper(t.value)
+            w, h = box(canvas, t.label, font)
+            if self.vertical:
+                draw_text(canvas, t.label, pos, p, font, "left" if sign > 0 else "right", "center")
+                widest = max(widest, w)
+            else:
+                draw_text(canvas, t.label, p, pos, font, "center", "top" if sign > 0 else "bottom")
+                widest = max(widest, h)
+        label = self.cb.auto_label()
+        if label is None:
+            return
+        lf = st.label_font()
+        pos += sign * (widest + axis.get("label_pad"))
+        if self.vertical:
+            draw_text(canvas, label, pos, (y0 + y1) / 2, lf, "center",
+                      "top" if sign > 0 else "bottom", 90.0)
+        else:
+            draw_text(canvas, label, (x0 + x1) / 2, pos, lf, "center",
+                      "top" if sign > 0 else "bottom")