plot-io 0.12.0__py3-none-win_arm64.whl

plot_io/presets/heatmap.py

@@ -0,0 +1,322 @@
+"""Heatmap: a grid of values drawn as colored cells (``Cells``), with optional cell labels
+(``Text``)."""
+
+import array
+import math
+from typing import Any
+
+from vec_io.arrays import doubles
+
+from .. import groups
+from .. import kernels
+from .._methods import HeatmapMixin
+from ..color.values import color, ink_for
+from ..layers.base import Series, common
+from ..layers.core import Call, Layer, Table
+from ..layers.geoms import GEOMS
+from ..layers.paint import Paint, role
+from ..numbers import NumberStyle, check_format, format_numbers
+from ..mapped import color_axis_of, trained
+from ..settings import Setting, boolean, one_of, table
+from ..labels import FIT, fits
+from ..text import font_of
+from ..transform import settings_both as coord_settings
+from ._util import columns_meta, gradient_entry, placed_text
+
+__all__ = ["Heatmap", "Coords"]
+
+
+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 axis(self, call: Call, role: str) -> None:
+        """What the call's axis ``role`` needs: categories, the extent, the coordinates
+        given (in the extent), integer ticks for indices."""
+        if self.categories is not None:
+            call.cats[role] = self.categories
+        call.extent_of[role] = self.extent
+        if self.kind in ("centers", "edges"):
+            call.extra[role] = self.values
+        elif self.kind == "index":
+            call.whole.add(role)
+
+    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 = 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])]
+
+    def extent(self, scale: Any) -> tuple[int, float, float]:
+        if self.kind in ("index", "category"):
+            return self.n, -0.5, self.n - 0.5
+        e = self.edges_f(scale.f)
+        lo, hi = min(e), max(e)
+        return self.n, scale.inverse(lo), scale.inverse(hi)
+
+
+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:]))
+
+
+class Heatmap(HeatmapMixin, Series):
+    """A grid of values (rows first) as colored cells; ``x``/``y`` give cell labels,
+    centres or edges.
+
+    A 3-D ``(h, w, 3|4)`` grid is drawn as colors.  ``x``/``y``: ``None`` (cell indices),
+    labels (a category axis), n cell centres or n + 1 edges (uneven spacing and any axis
+    scale allowed).  ``color_map``, ``color_limits`` and ``color_scale`` configure the
+    series' `color axis <docs/api.md#6-color-and-size-axes>`_.  ``value_labels=True`` writes each value, formatted by
+    ``value_label_format`` (any `number format <docs/api.md#number-formats>`_; ``"auto"``: 3
+    significant digits, the same decimals for all cells), in black or white, whichever
+    contrasts more (at most 2,500 cells unless ``value_labels="all"``).  ``kind``:
+    ``image``, ``cells`` or ``smooth``.
+
+    Example::
+
+        import math
+        p = plt.Plot(width=240, height=160)
+        p.heatmap([[math.sin(i / 2 + 0.5) * math.cos(j / 3) for j in range(6)] for i in range(5)],
+                  x=list("abcdef"), value_labels=True, colorbar=True)
+    """
+
+    SETTINGS = table(
+        **coord_settings(),
+        **common("color_map"),
+        kind=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."),
+        value_labels=Setting(_labels,
+                             doc="Write each cell's value (`\"all\"` above 2,500 cells)."),
+        value_label_format=Setting(check_format,
+                                   doc="Number format of cell labels (`auto`: 3 significant "
+                                       "digits, the same decimals for all cells)."),
+        **groups.TextStyle.embed("value_label_",
+                                 only=("font_size",),
+                                 docs={"font_size": "Font size of cell labels."}),
+        **common("name", "opacity", "z", "rasterize", "simplify", "font_size", "font_family",
+                 "text_color"),
+        **groups.NumberFormat.embed(),
+    )
+
+    DATA = {"values": "raw", "x": "raw", "y": "raw"}
+
+    def __init__(self, values: Any, x: Any = None, y: Any = None, plot: Any = None,
+                 **settings: Any) -> None:
+        values, x, y = self.cols["values"], self.cols["x"], self.cols["y"]
+        g = kernels.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 = 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")
+        super().__init__(plot=plot, **settings)
+
+    def _bar_source(self) -> Any:
+        return self if self.k == 1 else None          # RGB cells: no color axis
+
+    def _grid(self) -> tuple[Any, int, int, int, Coords, Coords]:
+        """``(flat values, rows, columns, channels, x coordinates, y coordinates)``."""
+        return self.flat, self.h, self.w, self.k, self.xc, self.yc
+
+    def _inverts(self) -> bool:
+        return self.resolve("origin") == "top"
+
+    def layers(self) -> Call:
+        call = Call(self)
+        flat, h, w, k, xc, yc = self._grid()
+        axis = color_axis_of(self) if self.k == 1 else None
+        cmap = axis.color_map() if axis is not None else None
+        key = id(self)
+
+        def page_edges(ctx: Any, c: Coords, 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 c.edges_f(state.f)]
+
+        def grid(ctx: Any) -> dict[str, Any]:
+            def make() -> dict[str, Any]:
+                xe = page_edges(ctx, xc, "x")
+                ye = page_edges(ctx, yc, "y")
+                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]
+                if axis is not None:
+                    st = trained(axis)
+                    assert st.scale is not None
+                    flo, fhi = st.flo, st.fhi
+                    fn = st.scale.function
+                    f = kernels.scaled(flat, 1.0, 0.0, fn, **st.scale.options) if fn else flat
+                    if fn and (math.inf in flat or -math.inf in flat):
+                        for i, v in enumerate(flat):        # infinities: the map's extremes
+                            if v in (math.inf, -math.inf):
+                                f[i] = v
+                    rows: list[Any] = [f[r * w:(r + 1) * w] for r in range(h)]
+                else:
+                    flo = fhi = 0.0
+                    rows = [[tuple(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]
+                return {"rows": rows, "xe": xe, "ye": ye, "flo": flo, "fhi": fhi}
+            return memo_grid(ctx, key, make)
+
+        if h and w:                                 # an empty grid draws nothing
+            call.layers.append(Layer(GEOMS["cells"], Table(),
+                                     paint=Paint(self, "heatmap.cells", grid=grid, cmap=cmap),
+                                     stat="identity", occupy="all", key=False))
+        if self.resolve("value_labels") and h and w:
+            call.layers.append(placed_text(self, key, lambda ctx: self._label_items(ctx, grid(ctx),
+                                                                              axis),
+                                           lambda ctx: font_of(
+                                               self, self.resolve("value_label_font_size")),
+                                           overlay=False))
+        name = self.resolve("name")
+        if name is not None:
+            call.entries.append(gradient_entry(
+                self, name, lambda: axis.color_map() if axis is not None else None))
+        if axis is not None:
+            call.color_sources = [(axis, flat)]
+        xc.axis(call, "x")
+        yc.axis(call, "y")
+        if self._inverts():
+            call.invert = {"y"}
+        call.exact = {"x", "y"}
+        columns_meta(call, {"values": "x"})
+        return call
+
+    def _label_items(self, ctx: Any, g: dict[str, Any], axis: Any) -> list[Any]:
+        _, h, w, _, _, _ = self._grid()
+        if h * w > 2500 and self.resolve("value_labels") != "all":
+            raise ValueError(f"{h * w} cells is too many for labels "
+                             "(use value_labels=\"all\" to force)")
+        fmt = self.resolve("value_label_format")
+        rows, xe, ye, flo, fhi = g["rows"], g["xe"], g["ye"], g["flo"], g["fhi"]
+        cmap = axis.color_map() if axis is not None else None
+        scale = trained(axis).scale if axis is not None else None
+        fn = scale.inverse if scale 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]
+        font = font_of(self, self.resolve("value_label_font_size"))
+        fit = FIT * font.size
+        out = []
+        for (r, c, v), text in zip(cells, texts):
+            if not fits(ctx.canvas, text, font, xe[c + 1] - xe[c], ye[r + 1] - ye[r], fit):
+                continue                            # a label that does not fit is left out
+            if cmap is None:
+                col = v
+            elif v != v:
+                col = cmap.bad
+            else:
+                col = cmap((v - flo) / ((fhi - flo) or 1.0))
+            out.append((text, (xe[c] + xe[c + 1]) / 2, (ye[r] + ye[r + 1]) / 2, ink_for(col),
+                        "center", "center"))
+        return out
+
+
+def memo_grid(ctx: Any, key: Any, make: Any) -> Any:
+    from ._util import memo
+    return memo(ctx, (key, "grid"), make)
+
+
+@role("heatmap.cells")
+def _cells_paint(s: Heatmap, ctx: Any, glyph: bool, grid: Any = None,
+                 cmap: Any = None) -> dict[str, Any]:
+    """The cells of ``grid(ctx)``: through the color map ``cmap``, or colors."""
+    g = grid(ctx)
+    opacity = s.resolve("opacity")
+    mode = s.resolve("kind")
+    p: dict[str, Any] = {"x_edges": g["xe"], "y_edges": g["ye"],
+                         "even": _even(g["xe"]) and _even(g["ye"]), "mode": mode,
+                         "opacity": 1.0 if opacity is None else opacity}
+    if cmap is not None:
+        flat = array.array("d")
+        for row in g["rows"]:
+            flat.extend(row)
+        p["fill"], p["palette"], p["bounds"] = cmap.cells(flat, g["flo"], g["fhi"], mode)
+        nan = ctx.color(s.resolve("nan_color"))
+        p["bad"] = nan if nan is not None else cmap.bad
+    else:
+        p["fill"] = [c for row in g["rows"] for c in row]
+    return p

plot_io/presets/hexbin.py

@@ -0,0 +1,198 @@
+"""Hexagonal 2-D histograms (stat ``hexbin``): counts of points in hexagons
+(``Polygons``, one hexagon repeated), colored by a map."""
+
+import math
+from typing import Any
+
+import vec_io
+
+from .._methods import HexbinMixin
+from .. import groups
+from .. import columns, kernels
+from ..color.values import with_alpha
+from ..layers import stats
+from ..layers.base import Series, common
+from ..layers.core import Call, Layer, Table
+from ..layers.grammar import entry
+from ..layers.paint import Paint, role
+from ..layers.geoms import GEOMS
+from ..mapped import color_axis_of, colors_of
+from ..settings import Setting, number, one_of, table
+from ..layers.pipeline import render_of
+from ..transform import settings_both as coord_settings
+from ._util import columns_meta
+
+__all__ = ["Hexbin"]
+
+
+_S3 = math.sqrt(3.0)
+
+
+def _frame_of(ctx: Any) -> tuple:
+    """The frame a lattice is laid out in: the axes' mappers relative to their top-left
+    corner, and their size."""
+    x0, y0, x1, y1 = ctx.rect
+    xm, ym = ctx.xmap, ctx.ymap
+    return (xm.function, xm.a, xm.b - x0, ym.function, ym.a, ym.b - y0, x1 - x0, y1 - y0)
+
+
+def _same(f: tuple, g: tuple) -> bool:
+    return f[0] == g[0] and f[3] == g[3] and all(
+        abs(a - b) <= 1e-6 * max(1.0, abs(a), abs(b))
+        for a, b in zip(f[1:3] + f[4:], g[1:3] + g[4:]))
+
+
+class Hexbin(HexbinMixin, Series):
+    """A hexagonal 2-D histogram of points ``(x, y)``: ``bins`` regular hexagons across
+    the axes width.  The lattice is laid out in page units for the axes the layout placed
+    (its domain is the data's, its bins the axes'), so the hexagons are regular and tile
+    exactly whatever the axes' aspect.
+
+    The counts depend on the axes, and the color axis on the counts: the layout counts again
+    when a round places the axes differently (`layout <docs/api.md#layout>`_).  Points outside the
+    axes are not counted.  Colored by ``color_map``; a colorbar source.
+
+    Example::
+
+        rng = random.Random(1)
+        x = [rng.gauss(0, 1) for _ in range(1000)]
+        p = plt.Plot(width=240, height=160)
+        p.hexbin(x, [v + rng.gauss(0, 0.7) for v in x], bins=16)
+    """
+
+    SETTINGS = table(
+        **coord_settings(),
+        bins=Setting(number(lo=1, integer=True), doc="Hexagons across the axes width."),
+        weights=Setting(lambda v: columns.numbers(v, "weights"), doc="A weight per point (sums)."),
+        stat=Setting(one_of("count", "probability", "percent"),
+                     doc="`count`, `probability` or `percent`."),
+        min_count=Setting(number(lo=0), doc="Hexagons below this count are not drawn."),
+        **common("color_map"),
+        **groups.Stroke.embed("edge_",
+                              only=("color", "line_width"),
+                              docs={"color": "Hexagon outline color; default none.",
+                                    "line_width": "Outline width."}),
+        **common("name", "opacity", "z", "rasterize", "simplify", "font_size", "font_family",
+                 "text_color"),
+    )
+
+    DATA = {"x": "numbers", "y": "numbers"}
+
+    def __init__(self, x: Any, y: Any, plot: Any = None, **settings: Any) -> None:
+        super().__init__(plot=plot, **settings)
+
+    def _bar_source(self) -> Any:
+        return self
+
+    def _nominal_frame(self) -> tuple:
+        """Outside a render: the lattice over the data range and the figure's nominal axes
+        size (``width`` × ``height``)."""
+        fig = self.plot.figure
+        w, h = fig.reference_size()
+        out: list[Any] = []
+        for axis, values, (p0, p1) in ((self.plot.x_axis, self.cols["x"], (0.0, w)),
+                                        (self.plot.y_axis, self.cols["y"], (h, 0.0))):
+            sc = stats.axis_scale(axis)
+            n, lo, hi = kernels.extent(values, sc.function)
+            if not n:
+                lo, hi = 0.0, 1.0
+            if hi == lo:
+                lo, hi = lo - 0.5, hi + 0.5
+            name = "auto_padding" if axis.get("padding") == "auto" else "padding"
+            px, frac = axis.resolve_linear(name)      # px + fraction of the range
+            span = abs(p1 - p0)
+            d = (hi - lo) * frac
+            if px and span > 2 * px:
+                d = (hi - lo + 2 * d) * px / (span - 2 * px) + d
+            lo, hi = lo - d, hi + d
+            a = (p1 - p0) / (hi - lo)
+            out += [sc.function, a, p0 - a * lo]
+        return (*out, w, h)
+
+    def _frame(self) -> tuple | None:
+        """The frame of the axes the layout placed (``None`` in a render before the first
+        placement), else, outside a render, the nominal one."""
+        r = render_of(self.plot)
+        if r is None:
+            return self._nominal_frame()
+        ctx = r.placed.get(id(self.plot))
+        return None if ctx is None else _frame_of(ctx)
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        """``{"x", "y": hexagon centres (page units, relative to the axes' top-left corner),
+        "values", "dx": hexagon width (px)}`` for the placed axes (see :meth:`_frame`)."""
+        frame = self._frame()
+        if frame is None:
+            raise ValueError(f"{self!r}: the hexagons are counted once the axes are placed")
+        return stats.compute(self, "hexbin", (), frame, kind=None)
+
+    def page_stats(self, ctx: Any) -> dict[str, Any]:
+        """The stats for the axes of ``ctx``: those of the placed axes, or (placed
+        differently) counted again; the colors keep the color axis trained before."""
+        f = _frame_of(ctx)
+        placed = self._frame()
+        if placed is not None and _same(f, placed):
+            return self.stats
+        return stats.compute(self, "hexbin", (), f, kind=None)
+
+    @staticmethod
+    def hexagon(dx: float) -> list[tuple[float, float]]:
+        """Vertices of the regular hexagon (pointy top) of width ``dx`` centred on 0."""
+        r = dx / _S3
+        return [(0.5 * dx, -0.5 * r), (0.5 * dx, 0.5 * r), (0.0, r), (-0.5 * dx, 0.5 * r),
+                (-0.5 * dx, -0.5 * r), (0.0, -r)]
+
+    def layers(self) -> Call:
+        call = Call(self)
+        axis = color_axis_of(self)
+
+        def page(ctx: Any) -> dict[str, Any]:
+            return self.page_stats(ctx)
+
+        t = Table()
+        t.lazy["x"] = lambda ctx: kernels.scaled(page(ctx)["x"], 1.0, ctx.rect[0])
+        t.lazy["y"] = lambda ctx: kernels.scaled(page(ctx)["y"], 1.0, ctx.rect[1])
+        t.lazy["size"] = lambda ctx: page(ctx)["dx"] * 1.002
+        t.lazy["half"] = lambda ctx: (page(ctx)["dx"] / 2, page(ctx)["dx"] / _S3)
+
+        def shape(ctx: Any) -> Any:
+            dx = page(ctx)["dx"]
+            path = vec_io.Path()
+            pts = [(x / dx, y / dx) for x, y in self.hexagon(dx)]
+            path.move_to(*pts[0])
+            for p in pts[1:]:
+                path.line_to(*p)
+            path.close()
+            return path
+        t.lazy["shape"] = shape
+
+        layer = Layer(GEOMS["polygons"], t, paint=Paint(self, "hexbin.cells"), stat="hexbin",
+                      glyph="hexagon")
+        call.layers.append(layer)
+        name = self.resolve("name")
+        if name is not None:
+            call.entries.append(entry(self, name, [layer], fill="symbol", marker=layer.paint))
+
+        def counts() -> Any:          # trained once the axes are placed
+            return None if self._frame() is None else self.stats["values"]
+        call.color_sources = [(axis, counts)]
+        call.extra = {"x": self.cols["x"], "y": self.cols["y"]}
+        call.resolution = _frame_of
+        columns_meta(call, {})
+        return call
+
+
+@role("hexbin.cells")
+def _cells_paint(s: Hexbin, ctx: Any, glyph: bool) -> dict[str, Any]:
+    """The hexagons in their counts' colors (the legend's: the map at 0.6)."""
+    axis = color_axis_of(s)
+    a = s.resolve("opacity")
+    if glyph:
+        return {"fill": with_alpha(axis.color_map()(0.6), a)}
+    lc = ctx.color(s.resolve("edge_color"))
+    lw = s.resolve("edge_width")
+    return {"fill": colors_of(axis, s.page_stats(ctx)["values"]),
+            "stroke": lc if lw and lc is not None else None,
+            "line_width": lw or None, "opacity": 1.0 if a is None else a,
+            "simplify": s.resolve("simplify")}