plot-io 0.1.0__py3-none-any.whl

plot_io/series/errors.py

@@ -0,0 +1,100 @@
+"""Error bars attached to scatter and line series."""
+
+import array
+import math
+from typing import Any
+
+from .. import _backend
+from ..settings import color, number
+from ..style import resolve_color
+
+__all__ = ["ErrorBars", "error_value"]
+
+
+class ErrorBars:
+    """Error bars: ``values`` is a value or sequence (symmetric) or a pair
+    ``(low, high)``; ``color`` defaults to the series color, ``cap`` is the cap width
+    (pt; 0: no caps)."""
+
+    def __init__(self, values: Any, color: Any = None, line_width: float | None = None,
+                 cap: float = 3.0) -> None:
+        self.values = values
+        self.color = None if color is None else _color(color)
+        self.line_width = None if line_width is None else number(lo=0)(line_width)
+        self.cap = number(lo=0)(cap)
+        self._low: Any = None
+        self._high: Any = None
+
+    def __repr__(self) -> str:
+        return "ErrorBars()"
+
+    def bind(self, n: int) -> None:
+        """Convert the values for ``n`` points."""
+        v = self.values
+        if isinstance(v, (tuple, list)) and len(v) == 2 and not _is_number(v[0]) \
+                and hasattr(v[0], "__len__"):
+            low, high = _column(v[0], n), _column(v[1], n)
+        else:
+            low = high = _column(v, n)
+        self._low, self._high = low, high
+
+    def ends(self, center: Any) -> tuple[array.array, array.array]:
+        lo = array.array("d", (c - e for c, e in zip(center, self._low)))
+        hi = array.array("d", (c + e for c, e in zip(center, self._high)))
+        return lo, hi
+
+    def draw(self, ctx: Any, xs: Any, ys: Any, horizontal: bool, base_color: Any,
+             width: float, simplify: Any, opacity: float) -> None:
+        """Bars along y (or x when ``horizontal``) from data ``xs``, ``ys``."""
+        center = xs if horizontal else ys
+        lo, hi = self.ends(center)
+        c = resolve_color(self.color, ctx.palette) if self.color is not None else base_color
+        w = self.line_width if self.line_width is not None else width
+        if c is None or not w:
+            return
+        canvas = ctx.canvas
+        if horizontal:
+            px0, px1, py = ctx.xmap.many(lo), ctx.xmap.many(hi), ctx.ymap.many(ys)
+            canvas.draw_segments(px0, py, px1, py, color=c, width=w, opacity=opacity,
+                                 simplify=simplify)
+            if self.cap:
+                h = self.cap / 2
+                for px in (px0, px1):
+                    canvas.draw_segments(px, _backend.scaled(py, 1.0, -h), px,
+                                         _backend.scaled(py, 1.0, h), color=c, width=w,
+                                         opacity=opacity, simplify=simplify)
+        else:
+            px, py0, py1 = ctx.xmap.many(xs), ctx.ymap.many(lo), ctx.ymap.many(hi)
+            canvas.draw_segments(px, py0, px, py1, color=c, width=w, opacity=opacity,
+                                 simplify=simplify)
+            if self.cap:
+                h = self.cap / 2
+                for py in (py0, py1):
+                    canvas.draw_segments(_backend.scaled(px, 1.0, -h), py,
+                                         _backend.scaled(px, 1.0, h), py, color=c, width=w,
+                                         opacity=opacity, simplify=simplify)
+
+
+def _is_number(v: Any) -> bool:
+    return isinstance(v, (int, float)) and not isinstance(v, bool)
+
+
+def _color(v: Any) -> Any:
+    return color(v)
+
+
+def _column(v: Any, n: int) -> array.array:
+    if _is_number(v):
+        return array.array("d", [float(v)] * n)
+    col = _backend.doubles(v)
+    if len(col) != n:
+        raise ValueError(f"{len(col)} error values for {n} points")
+    return col
+
+
+def error_value(v: Any) -> ErrorBars:
+    if isinstance(v, ErrorBars):
+        return v
+    if _is_number(v) and not (v >= 0 and math.isfinite(v)):
+        raise ValueError(f"expected a non-negative error, got {v!r}")
+    return ErrorBars(v)

plot_io/series/fill.py

@@ -0,0 +1,168 @@
+"""Fill: the area between two curves."""
+
+import array
+import math
+from typing import Any
+
+from pdf_io import arrays as pdf_arrays
+
+from .. import _backend
+from ..settings import Setting, color, number, one_of, table
+from ..style import with_alpha
+from . import Series, common
+from .curves import curve_settings, expand, interpolation_of, line_args
+
+__all__ = ["Fill"]
+
+
+def _curve(v: Any, n: int, what: str) -> Any:
+    if isinstance(v, (int, float)) and not isinstance(v, bool):
+        return float(v)
+    c = _backend.doubles(v)
+    if len(c) != n:
+        raise ValueError(f"{what} has {len(c)} values for {n} points")
+    return c
+
+
+class Fill(Series):
+    """The area between ``y1`` and ``y2`` (a sequence or a constant) along ``x``;
+    ``orientation="horizontal"`` fills between x curves along y (``fill(y, x1, x2)``)."""
+
+    SETTINGS = table(
+        color=Setting(color, doc="Fill color; default: the next palette color."),
+        alpha=Setting(number(0, 1),
+                      doc="Opacity; default 0.3."),
+        line_color=Setting(color, doc="Outline color; default none."),
+        line_width=Setting(number(lo=0), inherited=True, doc="Outline width (pt).", length=True),
+        **curve_settings(),
+        orientation=Setting(one_of("vertical", "horizontal"),
+                            doc="`vertical`: between y curves; `horizontal`: between x "
+                                "curves."),
+        **common("name", "z", "rasterize", "simplify", "font_size", "font_family",
+                 "text_color"),
+    )
+
+    def __init__(self, x: Any, y1: Any, y2: Any = 0.0, plot: Any = None,
+                 **settings: Any) -> None:
+        self.t = _backend.doubles(x)
+        n = len(self.t)
+        rows = pdf_arrays.rows(y1)
+        if rows is None and isinstance(y1, (list, tuple)) and len(y1) == 2 and all(
+                hasattr(r, "__len__") and not isinstance(r, str) for r in y1):
+            rows = [_backend.doubles(r) for r in y1]
+        if rows is not None and len(rows) == 2 and (isinstance(y2, (int, float)) and y2 == 0):
+            y1, y2 = rows
+        self.v1 = _curve(y1, n, "y1")
+        self.v2 = _curve(y2, n, "y2")
+        super().__init__(plot=plot, **settings)
+
+    def _horizontal(self) -> bool:
+        return self.get("orientation") == "horizontal"
+
+    def values(self, role: str) -> Any:
+        if (role == "x") != self._horizontal():
+            return self.t
+        out = array.array("d")
+        for v in (self.v1, self.v2):
+            if isinstance(v, float):
+                out.append(v)
+            else:
+                out.extend(v)
+        return out
+
+    def occupy(self, ctx: Any, occ: Any) -> None:
+        from ..occupancy import MAX_VERTICES, thin
+        n = len(self.t)
+        v1, v2 = ([v] * n if isinstance(v, float) else v for v in (self.v1, self.v2))
+        t, v1, v2 = thin(self.t, v1, v2, limit=MAX_VERTICES)
+        if self._horizontal():
+            occ.band(ctx.ymap.many(t), ctx.xmap.many(v1), ctx.xmap.many(v2), True)
+        else:
+            occ.band(ctx.xmap.many(t), ctx.ymap.many(v1), ctx.ymap.many(v2))
+
+    def legend_points(self) -> tuple[Any, Any] | None:
+        return None
+
+    def draw(self, ctx: Any) -> None:
+        c = self.base_color(ctx)
+        alpha = self.get("alpha")
+        stroke = ctx.color(self.get("line_color"))
+        width = self.get("line_width")
+        canvas = ctx.canvas
+        args = line_args(self)
+        interp = interpolation_of(self)
+        missing = self.get("missing")
+        if not self._horizontal():
+            xs = ctx.xmap.many(self.t)
+            y1 = ctx.ymap.many(self.v1) if not isinstance(self.v1, float) else ctx.ymap(self.v1)
+            y2 = ctx.ymap.many(self.v2) if not isinstance(self.v2, float) else ctx.ymap(self.v2)
+            simplify = self.simplify_arg(lines=True)
+            if c is not None:
+                canvas.fill_between(xs, y1, y2, color=c, opacity=alpha, stroke=None,
+                                    simplify=simplify, **args)
+            if stroke is not None and width:
+                for y in (y1, y2):
+                    if isinstance(y, float):
+                        if len(xs):
+                            canvas.draw_lines([min(xs), max(xs)], [y, y], color=stroke,
+                                              width=width)
+                    else:
+                        canvas.draw_lines(xs, y, color=stroke, width=width, **args)
+            return
+        # Horizontal: the curves run along y; steps and skipped points are expanded here,
+        # curves drawn straight.
+        ys = ctx.ymap.many(self.t)
+        n = len(ys)
+        x1 = (ctx.xmap.many(self.v1) if not isinstance(self.v1, float)
+              else array.array("d", [ctx.xmap(self.v1)] * n))
+        x2 = (ctx.xmap.many(self.v2) if not isinstance(self.v2, float)
+              else array.array("d", [ctx.xmap(self.v2)] * n))
+        ys, (x1, x2) = expand(ys, [x1, x2], interp if interp.startswith("step") else "linear",
+                              missing)
+        n = len(ys)
+        runs: list[list[int]] = []
+        cur: list[int] = []
+        for i in range(n):
+            if all(math.isfinite(v) for v in (ys[i], x1[i], x2[i])):
+                cur.append(i)
+            elif cur:
+                runs.append(cur)
+                cur = []
+        if cur:
+            runs.append(cur)
+        canvas.save()
+        if c is not None:
+            canvas.fill_style = with_alpha(c, alpha)
+            for run in runs:
+                canvas.begin_path()
+                canvas.move_to(x1[run[0]], ys[run[0]])
+                for i in run[1:]:
+                    canvas.line_to(x1[i], ys[i])
+                for i in reversed(run):
+                    canvas.line_to(x2[i], ys[i])
+                canvas.close_path()
+                canvas.fill()
+        canvas.restore()
+        if stroke is not None and width:
+            canvas.draw_lines(x1, ys, color=stroke, width=width)
+            canvas.draw_lines(x2, ys, color=stroke, width=width)
+
+    def zero_default(self, role: str) -> bool:
+        return role == ("x" if self.get("orientation") == "horizontal" else "y")
+
+    def legend_colors(self, ctx: Any) -> tuple[Any, Any]:
+        c = self.base_color(ctx)
+        stroke = ctx.color(self.get("line_color")) if self.get("line_width") else None
+        return (with_alpha(c, self.get("alpha")) if c is not None else None), stroke
+
+    def legend_glyph(self, ctx: Any, x: float, y: float, w: float, h: float) -> None:
+        c = self.base_color(ctx)
+        canvas = ctx.canvas
+        if c is not None:
+            canvas.fill_style = with_alpha(c, self.get("alpha"))
+            canvas.fill_rect(x, y, w, h)
+        stroke = ctx.color(self.get("line_color"))
+        if stroke is not None and self.get("line_width"):
+            canvas.stroke_style = stroke
+            canvas.line_width = self.get("line_width")
+            canvas.stroke_rect(x, y, w, h)

plot_io/series/heatmap.py

@@ -0,0 +1,304 @@
+"""Heatmap: a grid of values drawn as colored cells."""
+
+import array
+import math
+from typing import Any
+
+from pdf_io import arrays as pdf_arrays
+
+from .. import _backend
+from ..maps import ColorMap, check_cmap
+from ..numbers import NumberStyle, check_format, format_numbers, number_settings
+from ..limits import check_limits
+from ..scales import NUMERIC_SCALES, get_scale
+from ..settings import Setting, boolean, color, either, literal, number, one_of, table
+from ..text import Font, draw_text
+from . import Series, common
+
+__all__ = ["Heatmap"]
+
+
+def _labels(v: Any) -> Any:
+    if v == "all":
+        return "all"
+    return boolean(v)
+
+
+class _Coords:
+    """Cell coordinates along one axis: categories, centres or edges."""
+
+    def __init__(self, given: Any, n: int, what: str) -> None:
+        self.n = n
+        self.categories: list[Any] | None = None
+        self.kind = "index"
+        self.values: array.array = array.array("d")
+        if given is None:
+            return
+        items = list(given)
+        if items and all(isinstance(v, str) for v in items):
+            if len(items) != n:
+                raise ValueError(f"{len(items)} {what} labels for {n} cells")
+            self.categories = items
+            self.kind = "category"
+            return
+        vals = _backend.doubles(given)
+        if len(vals) == n:
+            self.kind = "centers"
+        elif len(vals) == n + 1:
+            self.kind = "edges"
+        else:
+            raise ValueError(f"{what} needs {n} centres or {n + 1} edges, got {len(vals)}")
+        if any(not math.isfinite(v) for v in vals):
+            raise ValueError(f"{what} coordinates must be finite")
+        d = [b - a for a, b in zip(vals, vals[1:])]
+        if d and not (all(x > 0 for x in d) or all(x < 0 for x in d)):
+            raise ValueError(f"{what} coordinates must be monotonic")
+        self.values = vals
+
+    def edges_f(self, f: Any) -> list[float]:
+        """Cell edges after the scale function ``f``."""
+        if self.kind in ("index", "category"):
+            return [i - 0.5 for i in range(self.n + 1)]
+        fv = [f(v) for v in self.values]
+        if any(not math.isfinite(v) for v in fv):
+            raise ValueError("heatmap coordinates outside the axis scale's domain")
+        if self.kind == "edges":
+            return fv
+        if len(fv) == 1:
+            return [fv[0] - 0.5, fv[0] + 0.5]
+        mids = [(a + b) / 2 for a, b in zip(fv, fv[1:])]
+        return [fv[0] - (mids[0] - fv[0])] + mids + [fv[-1] + (fv[-1] - mids[-1])]
+
+
+class Heatmap(Series):
+    """A grid of values (rows first) as colored cells; ``x``/``y`` give cell labels,
+    centres or edges."""
+
+    SETTINGS = table(
+        color_map=Setting(check_cmap, inherited=True,
+                     doc="A `ColorMap`, a map name or a list of colors."),
+        color_limits=Setting(check_limits, doc="Override of the map's limits for this series."),
+        color_scale=Setting(one_of(*NUMERIC_SCALES), doc="Override of the map's scale."),
+        mode=Setting(one_of("image", "cells", "smooth"),
+                     doc="`image` (one image pixel per cell), `cells` (rectangles) or "
+                         "`smooth` (interpolated)."),
+        nan_color=Setting(color, doc="Color of missing cells; default: the map's `bad` "
+                                     "(transparent)."),
+        origin=Setting(one_of("top", "bottom"),
+                       doc="`top`: row 0 at the top (y axis inverted); `bottom`: at the bottom."),
+        labels=Setting(_labels, doc="Write each cell's value (`\"all\"` above 2,500 cells)."),
+        label_format=Setting(check_format,
+                             doc="Number format of cell labels (`auto`: 3 significant digits, "
+                                 "the same decimals for all cells)."),
+        label_size=Setting(number(positive=True), doc="Font size of cell labels (pt).", length=True),
+        colorbar=Setting(either(literal(False), literal(True)),
+                         doc="Attach a colorbar to the plot."),
+        **common("name", "alpha", "z", "rasterize", "simplify", "font_size", "font_family",
+                 "text_color"),
+        **number_settings(),
+    )
+    USES_PALETTE = False
+
+    def __init__(self, values: Any, x: Any = None, y: Any = None, plot: Any = None,
+                 **settings: Any) -> None:
+        g = pdf_arrays.grid(values)
+        if g is not None:
+            h, w, k, flat = g
+            self.h, self.w, self.k = h, w, k
+            self.flat = flat
+        else:
+            rows = [list(r) for r in values]
+            if not rows:
+                raise ValueError("a heatmap needs at least one row")
+            w = len(rows[0])
+            if any(len(r) != w for r in rows):
+                raise ValueError("heatmap rows must have the same length")
+            if w and rows[0] and isinstance(rows[0][0], (list, tuple)):
+                self.k = len(rows[0][0])
+                self.flat = array.array("d", [float(c) for r in rows for cell in r for c in cell])
+            else:
+                self.k = 1
+                self.flat = _backend.doubles([v for r in rows for v in r])
+            self.h, self.w = len(rows), w
+        if self.k not in (1, 3, 4):
+            raise ValueError("cells must be values or (r, g, b[, a]) colors")
+        self.xc = _Coords(x, self.w, "x")
+        self.yc = _Coords(y, self.h, "y")
+        self._map: Any = None
+        self._map_key: Any = None
+        self._bar: Any = None
+        super().__init__(plot=plot, **settings)
+        if self.get("colorbar"):
+            self._attach_bar()
+
+    def set(self, **settings: Any) -> "Heatmap":
+        super().set(**settings)
+        if settings.get("colorbar") and getattr(self, "plot", None) is not None:
+            self._attach_bar()
+        return self
+
+    def _attach_bar(self) -> None:
+        if self._bar is None and self.k == 1:
+            self._bar = self.plot.colorbar(self)
+
+    # -- color map ----------------------------------------------------------------------
+
+    def color_map(self) -> ColorMap | None:
+        """The map used: ``cmap`` with ``color_limits``/``color_scale`` applied."""
+        if self.k != 1:
+            return None
+        cmap, lim, sc = self.get("color_map"), self.get("color_limits"), self.get("color_scale")
+        key = (id(cmap) if isinstance(cmap, ColorMap) else repr(cmap), lim, sc)
+        if self._map is None or self._map_key != key:
+            if isinstance(cmap, ColorMap) and lim is None and sc is None:
+                self._map = cmap
+            else:
+                base = cmap if isinstance(cmap, ColorMap) else ColorMap(cmap)
+                self._map = ColorMap(base, limits=lim if lim is not None else base.limits,
+                                     scale=sc or base.scale)
+            self._map_key = key
+        return self._map
+
+    def color_sources(self) -> list[tuple[Any, Any]]:
+        m = self.color_map()
+        return [] if m is None else [(m, self.flat)]
+
+    def has_own_color(self) -> bool:
+        return True
+
+    # -- axes ---------------------------------------------------------------------------
+
+    def _coords(self, role: str) -> _Coords:
+        return self.xc if role == "x" else self.yc
+
+    def categories(self, role: str) -> list[Any] | None:
+        return self._coords(role).categories
+
+    def inverts(self, role: str) -> bool:
+        return role == "y" and self.get("origin") == "top"
+
+    def extent(self, role: str, scale: Any) -> tuple[int, float, float]:
+        c = self._coords(role)
+        if c.kind in ("index", "category"):
+            return c.n, -0.5, c.n - 0.5
+        e = c.edges_f(scale.f)
+        lo, hi = min(e), max(e)
+        return c.n, scale.inverse(lo), scale.inverse(hi)
+
+    def values(self, role: str) -> Any:
+        c = self._coords(role)
+        return c.values if c.kind in ("centers", "edges") else None
+
+    def zero_padding(self, role: str) -> bool:
+        return True
+
+    # -- drawing ------------------------------------------------------------------------
+
+    def _page_edges(self, ctx: Any, role: str) -> list[float]:
+        state = ctx.x if role == "x" else ctx.y
+        mapper = ctx.xmap if role == "x" else ctx.ymap
+        return [mapper.position(e) for e in self._coords(role).edges_f(state.f)]
+
+    def occupy(self, ctx: Any, occ: Any) -> None:
+        self.occupy_all(ctx, occ)
+
+    def draw(self, ctx: Any) -> None:
+        xe = self._page_edges(ctx, "x")
+        ye = self._page_edges(ctx, "y")
+        h, w, k = self.h, self.w, self.k
+        flip_rows = ye[0] > ye[-1]           # row 0 must come first from the top
+        flip_cols = xe[0] > xe[-1]
+        if flip_rows:
+            ye = ye[::-1]
+        if flip_cols:
+            xe = xe[::-1]
+        cmap = self.color_map()
+        alpha = self.get("alpha")
+        alpha = 1.0 if alpha is None else alpha
+        if cmap is not None:
+            flo, fhi = ctx.maps[id(cmap)][:2]
+            fn = get_scale(cmap.scale).function
+            flat = _backend.scaled(self.flat, 1.0, 0.0, fn) if fn else self.flat
+            rows = [flat[r * w:(r + 1) * w] for r in range(h)]
+        else:
+            flo = fhi = 0.0
+            rows = [[tuple(self.flat[(r * w + c) * k:(r * w + c + 1) * k]) for c in range(w)]
+                    for r in range(h)]
+        order = list(range(h))[::-1] if flip_rows else list(range(h))
+        rows = [rows[i] for i in order]
+        if flip_cols:
+            rows = [r[::-1] for r in rows]
+        even_x = _even(xe)
+        even_y = _even(ye)
+        mode = self.get("mode")
+        nan_color = ctx.color(self.get("nan_color"))
+        pdf_map = cmap.to_pdf_io() if cmap is not None else None
+        args = dict(cmap=pdf_map, vmin=flo if cmap is not None else None,
+                    vmax=fhi if cmap is not None else None, nan_color=nan_color,
+                    opacity=alpha)
+        canvas = ctx.canvas
+        if pdf_map is None:
+            args.pop("cmap")
+        if even_x and even_y:
+            canvas.draw_heatmap(rows, xe[0], ye[0], xe[-1] - xe[0], ye[-1] - ye[0], mode=mode,
+                                **args)
+        else:
+            if mode == "image":
+                mode = "cells"
+            canvas.draw_heatmap(rows, x_edges=xe, y_edges=ye, mode=mode, **args)
+        if self.get("labels"):
+            self._labels(ctx, rows, xe, ye, cmap, flo, fhi)
+
+    def _cell_color(self, v: Any, cmap: Any, flo: float, fhi: float) -> Any:
+        if cmap is None:
+            return v
+        if v != v:
+            return cmap.bad
+        span = (fhi - flo) or 1.0
+        return cmap((v - flo) / span)
+
+    def _labels(self, ctx: Any, rows: list[Any], xe: list[float], ye: list[float], cmap: Any,
+                flo: float, fhi: float) -> None:
+        if self.h * self.w > 2500 and self.get("labels") != "all":
+            raise ValueError(f"{self.h * self.w} cells is too many for labels "
+                             "(use labels=\"all\" to force)")
+        fmt = self.get("label_format")
+        size = self.get("label_size")
+        family = self.get("font_family")
+        fn = get_scale(cmap.scale).inverse if cmap is not None else None
+        cells = [(r, c, v) for r, row in enumerate(rows) for c, v in enumerate(row)]
+        if cmap is not None:
+            cells = [(r, c, v) for r, c, v in cells if v == v]
+            style = NumberStyle.of(self)
+            if style.precision is None and fmt == "auto":
+                style.precision = 3
+            texts = format_numbers([fn(v) if fn else v for _, _, v in cells], fmt, style)
+        elif callable(fmt):
+            texts = [fmt(v) for _, _, v in cells]
+        elif fmt not in ("auto", "percent", "sci", "eng", "si", "int"):
+            texts = [fmt.format(*v[:3]) for _, _, v in cells]
+        else:
+            texts = ["" for _ in cells]
+        for (r, c, v), text in zip(cells, texts):
+            col = self._cell_color(v, cmap, flo, fhi)
+            ink = "black" if _luminance(col) > 0.45 else "white"
+            draw_text(ctx.canvas, text, (xe[c] + xe[c + 1]) / 2, (ye[r] + ye[r + 1]) / 2,
+                      Font(size, family, color=ink), "center", "center")
+
+
+def _even(e: list[float]) -> bool:
+    if len(e) < 3:
+        return True
+    d0 = (e[-1] - e[0]) / (len(e) - 1)
+    return all(abs((b - a) - d0) <= 1e-6 * max(1.0, abs(d0)) for a, b in zip(e, e[1:]))
+
+
+def _luminance(c: Any) -> float:
+    if c is None:
+        return 1.0
+    from pdf_io.colors import to_rgba
+    rgba = to_rgba(c) if not isinstance(c, tuple) else c
+    if rgba is None:
+        return 1.0
+    r, g, b = (float(x) for x in rgba[:3])
+    return 0.2126 * r + 0.7152 * g + 0.0722 * b