plot-io 0.1.0__py3-none-any.whl

plot_io/series/area.py

@@ -0,0 +1,181 @@
+"""Stacked areas: layers piled on a baseline (zero, symmetric or streamgraph)."""
+
+import array
+import math
+from typing import Any
+
+from pdf_io import arrays as pdf_arrays
+
+from .. import _backend
+from ..settings import Setting, boolean, color, number, one_of, table
+from ..style import with_alpha
+from . import Series, common
+from .curves import curve_settings, line_args
+
+__all__ = ["Area"]
+
+
+def _layers(ys: Any) -> tuple[list[Any], list[array.array]]:
+    if isinstance(ys, dict):
+        return list(ys.keys()), [_backend.doubles(v) for v in ys.values()]
+    rows = pdf_arrays.rows(ys)
+    if rows is None:
+        rows = [_backend.doubles(r) for r in ys]
+    return list(range(len(rows))), rows
+
+
+class Area(Series):
+    """Areas of the layers ``ys`` (rows of a 2-D sequence, or a dict ``{name: values}``)
+    along ``x``, stacked on ``baseline``."""
+
+    SETTINGS = table(
+        baseline=Setting(one_of("zero", "sym", "wiggle", "weighted_wiggle"),
+                         doc="`zero`, `sym` (symmetric around 0), `wiggle` or "
+                             "`weighted_wiggle` (streamgraphs, as matplotlib's stackplot)."),
+        stack=Setting(boolean, doc="Stack the layers; `False` overlays them from 0."),
+        color=Setting(lambda v: [color(c) for c in v] if isinstance(v, (list, tuple))
+                      and not (len(v) in (3, 4) and all(isinstance(c, (int, float)) for c in v))
+                      else color(v),
+                      doc="A color, or one per layer; default palette colors."),
+        alpha=Setting(number(0, 1), doc="Opacity of the fills."),
+        line_color=Setting(color, doc="Outline color of each layer's top; default none."),
+        line_width=Setting(number(lo=0), inherited=True, doc="Outline width (pt).", length=True),
+        **curve_settings(),
+        **common("name", "z", "rasterize", "simplify", "font_size", "font_family", "text_color"),
+    )
+
+    def __init__(self, x: Any, ys: Any, plot: Any = None, **settings: Any) -> None:
+        self.x = _backend.doubles(x)
+        self.names, self.ys = _layers(ys)
+        self.named = isinstance(ys, dict)
+        for r in self.ys:
+            if len(r) != len(self.x):
+                raise ValueError(f"x has {len(self.x)} values and a layer {len(r)}")
+        super().__init__(plot=plot, **settings)
+
+    def _horizontal(self) -> bool:
+        return False
+
+    @property
+    def stats(self) -> list[tuple[array.array, array.array]]:
+        """``(bottom, top)`` of each layer."""
+        return self.cached("stats", self._compute)
+
+    def _compute(self) -> list[tuple[array.array, array.array]]:
+        n = len(self.x)
+        ys = [[v if math.isfinite(v) else 0.0 for v in r] for r in self.ys]
+        if not self.get("stack"):
+            return [(array.array("d", bytes(8 * n)), array.array("d", r)) for r in self.ys]
+        m = len(ys)
+        total = [math.fsum(r[j] for r in ys) for j in range(n)]
+        base = self.get("baseline")
+        if base == "zero":
+            first = [0.0] * n
+        elif base == "sym":
+            first = [-t / 2 for t in total]
+        elif base == "wiggle":
+            first = [-math.fsum(ys[i][j] * (m - 0.5 - i) for i in range(m)) / m
+                     for j in range(n)]
+        else:
+            inv = [1.0 / t if t else 0.0 for t in total]
+            center = [0.0] * n
+            acc = 0.0
+            for j in range(n):
+                s = 0.0
+                cum = 0.0
+                for i in range(m):
+                    cum += ys[i][j]
+                    inc = ys[i][j] - (ys[i][j - 1] if j else 0.0) if j else ys[i][j]
+                    below = total[j] - cum + 0.5 * ys[i][j]
+                    move_up = 0.5 if j == 0 else below * inv[j]
+                    s += (move_up - 0.5) * inc
+                acc += s
+                center[j] = acc
+            first = [c - 0.5 * t for c, t in zip(center, total)]
+        out = []
+        cur = first
+        for r, raw in zip(ys, self.ys):
+            top = [a + b for a, b in zip(cur, r)]
+            bot = array.array("d", cur)
+            tp = array.array("d", (t if math.isfinite(v) else math.nan
+                                   for t, v in zip(top, raw)))
+            out.append((bot, tp))
+            cur = top
+        return out
+
+    def values(self, role: str) -> Any:
+        if role == "x":
+            return self.x
+        out = array.array("d")
+        for b, t in self.stats:
+            out.extend(b)
+            out.extend(t)
+        return out
+
+    def occupy(self, ctx: Any, occ: Any) -> None:
+        from ..occupancy import MAX_VERTICES, thin
+        for b, t in self.stats:
+            x, b, t = thin(self.x, b, t, limit=MAX_VERTICES)
+            occ.band(ctx.xmap.many(x), ctx.ymap.many(b), ctx.ymap.many(t))
+
+    def legend_points(self) -> tuple[Any, Any] | None:
+        xs, ys = array.array("d"), array.array("d")
+        for b, t in self.stats:
+            xs.extend(self.x)
+            ys.extend(t)
+            xs.extend(self.x)
+            ys.extend(b)
+        return xs, ys
+
+    def zero_default(self, role: str) -> bool:
+        return role == "y" and self.get("baseline") == "zero"
+
+    def palette_count(self) -> int:
+        return len(self.ys)
+
+    def has_own_color(self) -> bool:
+        return self.is_set("color")
+
+    def legend_labels(self) -> list[Any]:
+        if self.named:
+            return [str(k) for k in self.names]
+        name = self.get("name")
+        return [] if name is None else [name]
+
+    def _colors(self, ctx: Any) -> list[Any]:
+        c = self.get("color")
+        n = len(self.ys)
+        if isinstance(c, list):
+            if len(c) != n:
+                raise ValueError(f"{len(c)} colors for {n} layers")
+            return [ctx.color(x) for x in c]
+        if c is not None:
+            return [ctx.color(c)] * n
+        return self.auto_colors(ctx)
+
+    def draw(self, ctx: Any) -> None:
+        xs = ctx.xmap.many(self.x)
+        args = line_args(self)
+        lc = ctx.color(self.get("line_color"))
+        lw = self.get("line_width")
+        for (b, t), col in zip(self.stats, self._colors(ctx)):
+            pb, pt = ctx.ymap.many(b), ctx.ymap.many(t)
+            if col is not None:
+                ctx.canvas.fill_between(xs, pt, pb, color=col, opacity=self.get("alpha"),
+                                        simplify=self.simplify_arg(lines=True), **args)
+            if lc is not None and lw:
+                ctx.canvas.draw_lines(xs, pt, color=lc, width=lw, **args)
+
+    def legend_colors_at(self, ctx: Any, i: int) -> tuple[Any, Any]:
+        cols = self._colors(ctx)
+        c = cols[i] if i < len(cols) else None
+        return (with_alpha(c, self.get("alpha")) if c is not None else None), None
+
+    def legend_glyph_at(self, ctx: Any, i: int, x: float, y: float, w: float, h: float) -> None:
+        fill, _ = self.legend_colors_at(ctx, i)
+        if fill is not None:
+            ctx.canvas.fill_style = fill
+            ctx.canvas.fill_rect(x, y, w, h)
+
+    def legend_glyph(self, ctx: Any, x: float, y: float, w: float, h: float) -> None:
+        self.legend_glyph_at(ctx, 0, x, y, w, h)

plot_io/series/bar.py

@@ -0,0 +1,294 @@
+"""Bars: one rectangle per category or x value, grouped or stacked."""
+
+import array
+import math
+from typing import Any
+
+from .. import _backend
+from ..labels import bar_labels, check_labels
+from ..numbers import NumberStyle, check_format, format_numbers, number_settings
+from ..scales import get_scale
+from ..settings import Setting, boolean, color, number, one_of, table
+from ..shapes import draw_rects
+from ..slots import slots
+from ..stacking import check_stack, stacked
+from ..style import desaturate, with_alpha
+from ..text import Font
+from . import Series, common
+from .errors import error_value
+from ..values import Values, category_paint
+
+__all__ = ["Bar", "category_items"]
+
+
+def category_items(x: Any, values: Any) -> tuple[list[Any], Any, bool]:
+    """``(keys, values, categorical)`` of bar-like input: ``x`` and values, a dict, or
+    values alone (positions 0, 1, …)."""
+    if values is None:
+        if isinstance(x, dict):
+            keys = list(x.keys())
+            return keys, _backend.doubles(list(x.values())), all(isinstance(k, str) for k in keys)
+        vals = _backend.doubles(x)
+        return [float(i) for i in range(len(vals))], vals, False
+    vals = _backend.doubles(values)
+    if isinstance(x, (list, tuple)) and x and all(isinstance(k, str) for k in x):
+        keys = list(x)
+        categorical = True
+    else:
+        keys = list(_backend.doubles(x))
+        categorical = False
+    if len(keys) != len(vals):
+        raise ValueError(f"{len(keys)} positions for {len(vals)} values")
+    return keys, vals, categorical
+
+
+def _bottom(v: Any) -> Any:
+    if isinstance(v, (int, float)) and not isinstance(v, bool):
+        return float(v)
+    return _backend.doubles(v)
+
+
+def _width(v: Any) -> float:
+    if hasattr(v, "total_seconds"):
+        v = v.total_seconds()
+    return number(positive=True)(v)
+
+
+class Bar(Series):
+    """Bars of ``height`` at ``x`` (categories or numbers), or of a dict
+    ``{category: height}``, or of heights alone (at 0, 1, …)."""
+
+    SETTINGS = table(
+        color=Setting(category_paint,
+                      doc="Bar fill: a color, one per bar (list or dict by category), or "
+                          "numbers per bar (list or dict) mapped through `color_map`; "
+                          "default: the next palette color."),
+        **common("color_map"),
+        line_color=Setting(color, doc="Outline color; default none."),
+        line_width=Setting(number(lo=0), doc="Outline width (pt).", length=True),
+        width=Setting(_width,
+                      doc="Bar width: a category's slot (category units) or data units (a "
+                          "`timedelta` on time axes)."),
+        bottom=Setting(_bottom, doc="Base of the bars: a value or one per bar."),
+        orientation=Setting(one_of("vertical", "horizontal"),
+                            doc="`vertical` (bars rise from the x axis) or `horizontal`."),
+        stack=Setting(check_stack,
+                      doc="`True`: on top of the previous stacked bars of the plot; "
+                          "`\"normalize\"`: stacks scaled to 1."),
+        dodge=Setting(boolean, doc="Bar series sharing a category share its slot side "
+                                         "by side; `False` overlays them."),
+        gap=Setting(number(0, 1), doc="Fraction of each bar's part of a shared slot "
+                                           "left empty."),
+        saturation=Setting(number(0, 1), doc="Saturation of the fills (seaborn: 0.75)."),
+        error=Setting(error_value, doc="Error bars at the bar ends: a value, a sequence, "
+                                       "`(low, high)` or `ErrorBars`."),
+        error_color=Setting(color, doc="Error bar color."),
+        labels=Setting(check_labels,
+                       doc="Value labels: `True`/`\"outside\"`, `\"inside\"`, `\"center\"`."),
+        label_format=Setting(check_format, doc="Number format of the labels."),
+        **common("name", "alpha", "z", "rasterize", "simplify", "font_size", "font_family",
+                 "text_color"),
+        **number_settings(),
+    )
+    DODGE_GROUP = "bar"
+    STACK_GROUP = "bar"
+
+    def __init__(self, x: Any, height: Any = None, plot: Any = None, **settings: Any) -> None:
+        keys, vals, categorical = category_items(x, height)
+        self.cats = keys
+        self.heights = vals
+        self.positions = None if categorical else [float(k) for k in keys]
+        super().__init__(plot=plot, **settings)
+
+    def column_role(self, param: str) -> str | None:
+        if param == "x":
+            return self._cat_role()
+        return super().column_role(param)
+
+    def _horizontal(self) -> bool:
+        return self.get("orientation") == "horizontal"
+
+    def _cat_role(self) -> str:
+        return "y" if self._horizontal() else "x"
+
+    def stack_layers(self) -> list[tuple[list[Any], Any]]:
+        return [(self.cats, self.heights)]
+
+    def _ends(self) -> tuple[list[float], list[float]]:
+        """Bottoms and tops of the bars (data units)."""
+        if self.get("stack"):
+            return stacked(self)[0]
+        b = self.get("bottom")
+        n = len(self.heights)
+        bottoms = [b] * n if isinstance(b, float) else list(b)
+        if len(bottoms) != n:
+            raise ValueError(f"{len(bottoms)} bottoms for {n} bars")
+        return bottoms, [lo + h for lo, h in zip(bottoms, self.heights)]
+
+    # -- axes -------------------------------------------------------------------------
+
+    def categories(self, role: str) -> list[Any] | None:
+        if role == self._cat_role() and self.positions is None:
+            return list(self.cats)
+        return None
+
+    def values(self, role: str) -> Any:
+        if role == self._cat_role():
+            return array.array("d", self.positions or [])
+        bottoms, tops = self._ends()
+        out = array.array("d", bottoms)
+        out.extend(tops)
+        e = self.get("error")
+        if e is not None:
+            e.bind(len(tops))
+            lo, hi = e.ends(tops)
+            out.extend(lo)
+            out.extend(hi)
+        return out
+
+    def extent(self, role: str, scale: Any) -> tuple[int, float, float]:
+        if role == self._cat_role() and self.positions is not None:
+            w = self.get("width") / 2
+            n, lo, hi = _backend.extent(self.positions, scale.function)
+            return n, lo - w, hi + w
+        return super().extent(role, scale)
+
+    def zero_default(self, role: str) -> bool:
+        return role != self._cat_role()
+
+    def zero_padding(self, role: str) -> bool:
+        return False
+
+    def _per_bar(self, v: Any) -> list[float]:
+        if isinstance(v, dict):
+            if self.cats is None:
+                raise ValueError("a dict of colors needs categorical bars")
+            return [v.get(k, float("nan")) for k in self.cats]
+        if len(v) != len(self.heights):
+            raise ValueError(f"{len(v)} values for {len(self.heights)} bars")
+        return list(v)
+
+    def color_sources(self) -> list[tuple[Any, Any]]:
+        c = self.get("color")
+        if isinstance(c, Values):
+            return [(self.color_map_obj(), array.array("d", self._per_bar(c.data)))]
+        return []
+
+    def has_own_color(self) -> bool:
+        return self.is_set("color")
+
+    def occupy(self, ctx: Any, occ: Any) -> None:
+        bottoms, tops = self._ends()
+        pos = self.positions or [float(i) for i in range(len(self.cats))]
+        w = self.get("width") / 2
+        lo = [p - w for p in pos]
+        hi = [p + w for p in pos]
+        if self._horizontal():
+            occ.rects(ctx.xmap.many(bottoms), ctx.ymap.many(lo), ctx.xmap.many(tops),
+                      ctx.ymap.many(hi))
+        else:
+            occ.rects(ctx.xmap.many(lo), ctx.ymap.many(bottoms), ctx.xmap.many(hi),
+                      ctx.ymap.many(tops))
+
+    def legend_points(self) -> tuple[Any, Any] | None:
+        bottoms, tops = self._ends()
+        pos = self.positions or [float(i) for i in range(len(self.cats))]
+        if self._horizontal():
+            return array.array("d", tops), array.array("d", pos)
+        return array.array("d", pos), array.array("d", tops)
+
+    # -- drawing ----------------------------------------------------------------------
+
+    def fills(self, ctx: Any) -> list[Any]:
+        c = self.get("color")
+        n = len(self.heights)
+        if isinstance(c, Values):
+            cm = self.color_map_obj()
+            flo, fhi = ctx.maps[id(cm)][:2]
+            f = get_scale(cm.scale).f
+            span = (fhi - flo) or 1.0
+            cols = [cm((f(v) - flo) / span) for v in self._per_bar(c.data)]
+        elif isinstance(c, dict):
+            if self.cats is None:
+                raise ValueError("a dict of colors needs categorical bars")
+            cols = [ctx.color(c[k]) if c.get(k) is not None else None for k in self.cats]
+        elif isinstance(c, list):
+            if len(c) != n:
+                raise ValueError(f"{len(c)} colors for {n} bars")
+            cols = [ctx.color(x) for x in c]
+        else:
+            cols = [self.base_color(ctx)] * n
+        sat = self.get("saturation")
+        alpha = self.get("alpha")
+        alpha = 1.0 if alpha is None else alpha
+        return [with_alpha(desaturate(x, sat) if sat != 1.0 else x, alpha)
+                if x is not None else None for x in cols]
+
+    def rects(self, ctx: Any) -> list[tuple[float, float, float, float, bool]]:
+        """Page rectangles ``(x0, y0, x1, y1, negative)`` of the bars."""
+        horiz = self._horizontal()
+        cmap, vmap = (ctx.ymap, ctx.xmap) if horiz else (ctx.xmap, ctx.ymap)
+        bottoms, tops = self._ends()
+        out = []
+        for (center, width), b, t in zip(slots(self, ctx), bottoms, tops):
+            a0, a1 = cmap(center - width / 2), cmap(center + width / 2)
+            v0, v1 = vmap(b), vmap(t)
+            neg = t < b
+            if horiz:
+                out.append((min(v0, v1), min(a0, a1), max(v0, v1), max(a0, a1), neg))
+            else:
+                out.append((min(a0, a1), min(v0, v1), max(a0, a1), max(v0, v1), neg))
+        return out
+
+    def draw(self, ctx: Any) -> None:
+        rects = self.rects(ctx)
+        fills = self.fills(ctx)
+        lc = ctx.color(self.get("line_color"))
+        lw = self.get("line_width")
+        keep = [i for i, r in enumerate(rects) if all(math.isfinite(v) for v in r[:4])]
+        draw_rects(ctx.canvas, [rects[i][0] for i in keep], [rects[i][1] for i in keep],
+                   [rects[i][2] - rects[i][0] for i in keep],
+                   [rects[i][3] - rects[i][1] for i in keep],
+                   fill=[fills[i] for i in keep], stroke=lc if lw else None,
+                   line_width=lw or None, simplify=self.simplify_arg())
+        e = self.get("error")
+        if e is not None:
+            bottoms, tops = self._ends()
+            e.bind(len(tops))
+            centers = [c for c, _ in slots(self, ctx)]
+            ec = ctx.color(self.get("error_color"))
+            if self._horizontal():
+                e.draw(ctx, tops, centers, True, ec, 0.8, None, 1.0)
+            else:
+                e.draw(ctx, centers, tops, False, ec, 0.8, None, 1.0)
+
+    def draw_overlay(self, ctx: Any) -> None:
+        where = self.get("labels")
+        if not where:
+            return
+        rects = self.rects(ctx)
+        fills = self.fills(ctx)
+        keep = [i for i, r in enumerate(rects) if all(math.isfinite(v) for v in r[:4])]
+        bottoms, tops = self._ends()
+        vals = [t - b for b, t in zip(bottoms, tops)]
+        texts = format_numbers(vals, self.get("label_format"), NumberStyle.of(self))
+        font = Font(self.get("font_size"), self.get("font_family"),
+                    color=ctx.color(self.get("text_color")))
+        bar_labels(ctx.canvas, [rects[i] for i in keep], [texts[i] for i in keep],
+                   [fills[i] for i in keep], where, self._horizontal(), font)
+
+    def legend_colors(self, ctx: Any) -> tuple[Any, Any]:
+        fills = self.fills(ctx)
+        lc = ctx.color(self.get("line_color")) if self.get("line_width") else None
+        return (fills[0] if fills else None), lc
+
+    def legend_glyph(self, ctx: Any, x: float, y: float, w: float, h: float) -> None:
+        fill, stroke = self.legend_colors(ctx)
+        canvas = ctx.canvas
+        canvas.begin_path()
+        canvas.rect(x, y, w, h)
+        if fill is not None:
+            canvas.fill(fill)
+        if stroke is not None:
+            canvas.line_width = self.get("line_width")
+            canvas.stroke(stroke)