plot-io 0.1.0__py3-none-any.whl

plot_io/series/hist2d.py

@@ -0,0 +1,163 @@
+"""2-D histograms: counts of points in rectangular bins, drawn as a heatmap."""
+
+import array
+import math
+from typing import Any
+
+from .. import _backend, stats
+from ..scales import get_scale
+from ..settings import Setting, either, literal, number, one_of, pair, table
+from .heatmap import Heatmap, _Coords
+
+__all__ = ["Hist2d"]
+
+
+def _bins2(v: Any) -> Any:
+    if isinstance(v, (tuple, list)) and len(v) == 2 and not all(
+            isinstance(x, (int, float)) for x in v):
+        return (stats.check_bins(v[0]), stats.check_bins(v[1]))
+    if isinstance(v, (tuple, list)) and len(v) == 2 and all(
+            isinstance(x, int) and not isinstance(x, bool) for x in v):
+        return (stats.check_bins(v[0]), stats.check_bins(v[1]))
+    b = stats.check_bins(v)
+    return (b, b)
+
+
+def _range2(v: Any) -> Any:
+    return (pair(number())(v[0]), pair(number())(v[1]))
+
+
+class Hist2d(Heatmap):
+    """A 2-D histogram of points ``(x, y)``: counts in bins (even on the axis scales),
+    drawn as a heatmap; bins with fewer than ``min_count`` points are left empty."""
+
+    SETTINGS = table(
+        **{k: v for k, v in Heatmap.SETTINGS.items() if k not in ("origin",)},
+        bins=Setting(_bins2,
+                     doc="Bins of both axes, or `(x bins, y bins)`: counts, edges or rules "
+                         "(as `Hist`)."),
+        range=Setting(either(literal(None), _range2),
+                      doc="`((x0, x1), (y0, y1))`; default the data range."),
+        weights=Setting(lambda v: _backend.doubles(v), doc="A weight per point."),
+        stat=Setting(one_of("count", "density", "probability", "percent"),
+                     doc="`count`, `density` (volume 1), `probability` or `percent`."),
+        min_count=Setting(number(lo=0), doc="Bins below this count are transparent."),
+        origin=Setting(one_of("top", "bottom"), doc="Row of the first y bin."),
+    )
+
+    def __init__(self, x: Any, y: Any, plot: Any = None, **settings: Any) -> None:
+        self.px, self.py = _backend.doubles(x), _backend.doubles(y)
+        if len(self.px) != len(self.py):
+            raise ValueError(f"x has {len(self.px)} values and y {len(self.py)}")
+        self.k = 1
+        self._map: Any = None
+        self._map_key: Any = None
+        self._bar: Any = None
+        from . import Series
+        Series.__init__(self, plot=plot, **settings)
+        if self.get("colorbar"):
+            self._attach_bar()
+
+    def inverts(self, role: str) -> bool:
+        return False
+
+    def _scales(self) -> tuple[Any, Any]:
+        out = []
+        for axis in (self.plot.x_axis, self.plot.y_axis):
+            name = axis.nominal_scale
+            out.append(get_scale("linear" if name == "category" else name))
+        return out[0], out[1]
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        """``{"x_edges", "y_edges", "counts": rows (first y bin first), "values": rows}``."""
+        sx, sy = self._scales()
+        return self.cached(("stats", sx.name, sy.name), lambda: self._compute(sx, sy))
+
+    def _compute(self, sx: Any, sy: Any) -> dict[str, Any]:
+        bx, by = self.get("bins")
+        rng = self.get("range")
+        fx = _backend.scaled(self.px, 1.0, 0.0, sx.function) if sx.function else self.px
+        fy = _backend.scaled(self.py, 1.0, 0.0, sy.function) if sy.function else self.py
+        w = self.get("weights")
+
+        def edges(b: Any, v: Any, s: Any, r: Any) -> Any:
+            if isinstance(b, array.array):
+                return array.array("d", (s.f(e) for e in b))
+            fr = None if r is None else (s.f(r[0]), s.f(r[1]))
+            return stats.bin_edges(v, b, fr)
+
+        xe = edges(bx, fx, sx, None if rng is None else rng[0])
+        ye = edges(by, fy, sy, None if rng is None else rng[1])
+        counts = stats.histogram2d(fx, fy, xe, ye, w)
+        total = math.fsum(math.fsum(r) for r in counts)
+        st = self.get("stat")
+        xd = [sx.inverse(e) for e in xe]
+        yd = [sy.inverse(e) for e in ye]
+        mc = self.get("min_count")
+        rows = []
+        for j, r in enumerate(counts):
+            out = array.array("d")
+            for i, c in enumerate(r):
+                if c < mc or c == 0:
+                    out.append(math.nan)
+                elif st == "count":
+                    out.append(c)
+                elif st == "probability":
+                    out.append(c / total)
+                elif st == "percent":
+                    out.append(100.0 * c / total)
+                else:
+                    area = (xe[i + 1] - xe[i]) * (ye[j + 1] - ye[j])
+                    out.append(c / (total * area) if area else math.nan)
+            rows.append(out)
+        return {"x_edges": array.array("d", xd), "y_edges": array.array("d", yd),
+                "counts": counts, "values": rows}
+
+    # the heatmap's grid, from the bins
+    @property
+    def flat(self) -> Any:
+        out = array.array("d")
+        for r in self.stats["values"]:
+            out.extend(r)
+        return out
+
+    @flat.setter
+    def flat(self, v: Any) -> None:
+        pass
+
+    @property
+    def h(self) -> int:
+        return len(self.stats["values"])
+
+    @h.setter
+    def h(self, v: Any) -> None:
+        pass
+
+    @property
+    def w(self) -> int:
+        rows = self.stats["values"]
+        return len(rows[0]) if rows else 0
+
+    @w.setter
+    def w(self, v: Any) -> None:
+        pass
+
+    @property
+    def xc(self) -> _Coords:
+        return _Coords(self.stats["x_edges"], self.w, "x")
+
+    @xc.setter
+    def xc(self, v: Any) -> None:
+        pass
+
+    @property
+    def yc(self) -> _Coords:
+        return _Coords(self.stats["y_edges"], self.h, "y")
+
+    @yc.setter
+    def yc(self, v: Any) -> None:
+        pass
+
+    def legend_points(self) -> tuple[Any, Any]:
+        return self.px, self.py

plot_io/series/kde.py

@@ -0,0 +1,228 @@
+"""Kernel density estimates: smooth curves of the distribution of values."""
+
+import array
+import math
+from typing import Any
+
+from .. import _backend, stats
+from ..scales import get_scale
+from ..settings import Setting, boolean, either, literal, number, one_of, pair, table
+from ..shapes import fill_polygons
+from ..stacking import check_stack, stacked
+from ..style import dash_of, with_alpha
+from . import Series, common
+from .hist import _weights, value_sets
+from .points import colors_only
+
+__all__ = ["Kde", "check_bandwidth"]
+
+
+def check_bandwidth(v: Any) -> Any:
+    if isinstance(v, str):
+        return one_of("scott", "silverman")(v)
+    return number(positive=True)(v)
+
+
+class Kde(Series):
+    """A Gaussian kernel density estimate of ``values`` (one array, or several as a dict or
+    a list of arrays, on a common grid)."""
+
+    SETTINGS = table(
+        bandwidth=Setting(check_bandwidth,
+                          doc="`scott`, `silverman`, or a bandwidth in data units."),
+        adjust=Setting(number(positive=True), doc="Factor applied to the bandwidth."),
+        cut=Setting(number(lo=0), doc="The curve extends this many bandwidths past the "
+                                           "data."),
+        clip=Setting(either(literal(None), pair(number())), doc="Limits of the curve."),
+        grid=Setting(number(lo=16, integer=True), doc="Number of points of the curve."),
+        cumulative=Setting(boolean, doc="Plot the cumulative distribution."),
+        fill=Setting(boolean, doc="Fill under the curve."),
+        weights=Setting(_weights, doc="A weight per value (per set for several sets)."),
+        orientation=Setting(one_of("vertical", "horizontal"),
+                            doc="`vertical` (values along x) or `horizontal`."),
+        stack=Setting(check_stack, doc="Stack the sets (seaborn's `multiple=\"stack\"`)."),
+        color=Setting(colors_only, doc="Color, or one per set; default palette colors."),
+        alpha=Setting(number(0, 1), doc="Fill opacity; default 0.25 (1 when stacked)."),
+        line_width=Setting(number(lo=0), inherited=True, doc="Line width (pt).", length=True),
+        line_style=Setting(one_of("solid", "dashed", "dotted", "dashdot"),
+                           doc="Line style."),
+        **common("name", "z", "rasterize", "simplify", "font_size", "font_family", "text_color"),
+    )
+    STACK_GROUP = "kde"
+
+    def __init__(self, values: Any, plot: Any = None, **settings: Any) -> None:
+        self.sets, self.named = value_sets(values)
+        super().__init__(plot=plot, **settings)
+
+    def _horizontal(self) -> bool:
+        return self.get("orientation") == "horizontal"
+
+    def _val_role(self) -> str:
+        return "y" if self._horizontal() else "x"
+
+    def _scale(self) -> Any:
+        axis = self.plot.y_axis if self._horizontal() else self.plot.x_axis
+        name = axis.nominal_scale
+        return get_scale("linear" if name == "category" else name)
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        """``{"grid": values, "densities": [per set], "bandwidths": [per set]}`` (densities
+        per unit of the axis scale on log axes)."""
+        scale = self._scale()
+        return self.cached(("stats", scale.name), lambda: self._compute(scale))
+
+    def _compute(self, scale: Any) -> dict[str, Any]:
+        fn = scale.function
+        weights = self.get("weights")
+        fsets, ws, bws = [], [], []
+        for i, (key, v) in enumerate(self.sets):
+            fv = _backend.scaled(v, 1.0, 0.0, fn) if fn else v
+            w = None
+            if weights is not None:
+                w = (weights[key] if isinstance(weights, dict) else
+                     weights[i] if isinstance(weights, list) else weights)
+            fsets.append(fv)
+            ws.append(w)
+            bws.append(stats.bandwidth(fv, self.get("bandwidth"), w) * self.get("adjust"))
+        cut = self.get("cut")
+        lo, hi = math.inf, -math.inf
+        for fv, bw in zip(fsets, bws):
+            n, a, b = _backend.extent(fv, None)
+            if n:
+                lo, hi = min(lo, a - cut * bw), max(hi, b + cut * bw)
+        if not math.isfinite(lo):
+            lo, hi = 0.0, 1.0
+        clip = self.get("clip")
+        if clip is not None:
+            lo, hi = max(lo, scale.f(clip[0])), min(hi, scale.f(clip[1]))
+        n = self.get("grid")
+        dens = [stats.kde(fv, lo, hi, n, bw, w, self.get("cumulative"))
+                for fv, bw, w in zip(fsets, bws, ws)]
+        grid = array.array("d", (scale.inverse(lo + (hi - lo) * i / (n - 1)) for i in range(n)))
+        return {"grid": grid, "densities": dens, "bandwidths": bws}
+
+    def stack_layers(self) -> list[tuple[list[Any], Any]]:
+        st = self.stats
+        keys = [round(g, 12) for g in st["grid"]]
+        return [(keys, d) for d in st["densities"]]
+
+    def _ends(self) -> list[tuple[list[float], list[float]]]:
+        if self.get("stack"):
+            return stacked(self)
+        return [([0.0] * len(d), list(d)) for d in self.stats["densities"]]
+
+    def values(self, role: str) -> Any:
+        if role == self._val_role():
+            return self.stats["grid"]
+        out = array.array("d", [0.0])
+        for _, t in self._ends():
+            out.extend(t)
+        return out
+
+    def zero_default(self, role: str) -> bool:
+        return role != self._val_role()
+
+    def zero_padding(self, role: str) -> bool:
+        return False
+
+    def palette_count(self) -> int:
+        return len(self.sets)
+
+    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.sets]
+        name = self.get("name")
+        return [] if name is None else [name]
+
+    def occupy(self, ctx: Any, occ: Any) -> None:
+        t = self.stats["grid"]
+        filled = self.get("fill") or self.get("stack")
+        for bottom, top in self._ends():
+            if self._horizontal():
+                if filled:
+                    occ.band(ctx.ymap.many(t), ctx.xmap.many(bottom), ctx.xmap.many(top), True)
+                else:
+                    occ.lines(ctx.xmap.many(top), ctx.ymap.many(t))
+            elif filled:
+                occ.band(ctx.xmap.many(t), ctx.ymap.many(bottom), ctx.ymap.many(top))
+            else:
+                occ.lines(ctx.xmap.many(t), ctx.ymap.many(top))
+
+    def legend_points(self) -> tuple[Any, Any] | None:
+        g = self.stats["grid"]
+        xs, ys = array.array("d"), array.array("d")
+        for _, t in self._ends():
+            xs.extend(g)
+            ys.extend(t)
+        return (ys, xs) if self._horizontal() else (xs, ys)
+
+    def colors(self, ctx: Any) -> list[Any]:
+        c = self.get("color")
+        n = len(self.sets)
+        if isinstance(c, list):
+            if len(c) != n:
+                raise ValueError(f"{len(c)} colors for {n} sets")
+            return [ctx.color(x) for x in c]
+        if c is not None:
+            return [ctx.color(c)] * n
+        return self.auto_colors(ctx)
+
+    def _fill_alpha(self) -> float:
+        return self.get("alpha")
+
+    def draw(self, ctx: Any) -> None:
+        g = self.stats["grid"]
+        horiz = self._horizontal()
+        vmap, dmap = (ctx.ymap, ctx.xmap) if horiz else (ctx.xmap, ctx.ymap)
+        pg = vmap.many(g)
+        lw = self.get("line_width")
+        dash = dash_of(self.get("line_style"), lw)
+        fill = self.get("fill") or self.get("stack")
+        for (bot, top), col in zip(self._ends(), self.colors(ctx)):
+            if col is None:
+                continue
+            pt, pb = dmap.many(top), dmap.many(bot)
+            if fill:
+                rx = array.array("d", pg)
+                rx.extend(reversed(pg))
+                ry = array.array("d", pt)
+                ry.extend(reversed(pb))
+                ring = (ry, rx) if horiz else (rx, ry)
+                fill_polygons(ctx.canvas, [[ring]], fill=with_alpha(col, self._fill_alpha()))
+            xs, ys = (pt, pg) if horiz else (pg, pt)
+            if lw:
+                ctx.canvas.draw_lines(xs, ys, color=col, width=lw, dash=dash,
+                                      simplify=self.simplify_arg(lines=True))
+
+    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
+        fill = with_alpha(c, self._fill_alpha()) if c is not None and (
+            self.get("fill") or self.get("stack")) else None
+        return fill, c
+
+    def legend_glyph_at(self, ctx: Any, i: int, x: float, y: float, w: float, h: float) -> None:
+        fill, c = self.legend_colors_at(ctx, i)
+        canvas = ctx.canvas
+        if fill is not None:
+            canvas.fill_style = fill
+            canvas.fill_rect(x, y, w, h)
+        if c is not None:
+            lw = self.get("line_width")
+            canvas.draw_lines([x, x + w], [y + h / 2, y + h / 2] if fill is None else [y, y],
+                              color=c, width=lw, dash=dash_of(self.get("line_style"), lw))
+
+    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)
+
+
+def _kde2d() -> type:
+    from .kde2d import Kde2d
+    return Kde2d
+
+
+Kde.with_y = staticmethod(_kde2d)          # type: ignore[attr-defined]

plot_io/series/kde2d.py

@@ -0,0 +1,153 @@
+"""2-D kernel density contours (``kde(x, y)``): seaborn's ``kdeplot(x=, y=)``."""
+
+import array
+import math
+from typing import Any
+
+from .. import _backend, stats
+from ..density import Density, _bandwidth
+from ..settings import Setting, color, number, table
+from .contour import Contour
+
+__all__ = ["Kde2d"]
+
+
+class Kde2d(Contour):
+    """Density contours of points ``(x, y)``: the 2-D Gaussian density on a grid over the
+    data plus ``cut`` bandwidths, drawn at levels enclosing equal shares of the mass."""
+
+    SETTINGS = table(
+        **{k: v for k, v in Contour.SETTINGS.items() if k not in ("levels", "colorbar")},
+        levels=Setting(lambda v: v if isinstance(v, list) else number(lo=1, integer=True)(v), doc="A number of iso-proportion levels (contour k of n encloses "
+                               "k/n of the mass, above `thresh`), or density values."),
+        thresh=Setting(number(0, 1),
+                       doc="Lowest iso-proportion: the outermost contour encloses 1 − "
+                           "`thresh` of the mass."),
+        bandwidth=Setting(_bandwidth,
+                          doc="`\"scott\"`, `\"silverman\"`, a number or a 2 × 2 matrix, as "
+                              "`Density`."),
+        adjust=Setting(number(positive=True), doc="Factor of the bandwidth."),
+        cut=Setting(number(lo=0), doc="Grid extends this many bandwidths past the data."),
+        gridsize=Setting(number(lo=16, hi=1024, integer=True),
+                         doc="Grid points per side."),
+        color=Setting(color, doc="Line color of unfilled contours without a `color_map`; "
+                                 "default the next palette color."),
+        colorbar=Setting(lambda v: bool(v), doc="Attach a colorbar to the plot."),
+    )
+
+    def __init__(self, x: Any, y: Any, plot: Any = None, **settings: Any) -> None:
+        self.px, self.py = _backend.doubles(x), _backend.doubles(y)
+        if len(self.px) != len(self.py):
+            raise ValueError(f"x has {len(self.px)} values and y {len(self.py)}")
+        self.h = self.w = 2
+        self.flat = array.array("d", [0.0] * 4)
+        self.xn = self.yn = array.array("d", [0.0, 1.0])
+        self._map = None
+        self._bar = None
+        from . import Series
+        Series.__init__(self, plot=plot, **settings)
+        if self.get("colorbar"):
+            self._bar = self.plot.colorbar(self)
+
+    @property
+    def USES_PALETTE(self) -> bool:                       # type: ignore[override]
+        return not self.get("filled") and not self.is_set("color_map")
+
+    def has_own_color(self) -> bool:
+        return self.is_set("color")
+
+    def _density(self) -> Density:
+        return Density(self.get("bandwidth"), self.get("adjust"))
+
+    def _grid(self) -> dict[str, Any]:
+        sx, sy = self._scales()
+        return self.cached(("grid", sx.name, sy.name), lambda: self._compute_grid(sx, sy))
+
+    def _compute_grid(self, sx: Any, sy: Any) -> dict[str, Any]:
+        fx = _backend.scaled(self.px, 1.0, 0.0, sx.function, **sx.options)
+        fy = _backend.scaled(self.py, 1.0, 0.0, sy.function, **sy.options)
+        d = self._density()
+        h = d.covariance(fx, fy)
+        n = int(self.get("gridsize"))
+        _, x0, x1 = _backend.extent(fx, None)
+        _, y0, y1 = _backend.extent(fy, None)
+        if not all(math.isfinite(v) for v in (x0, x1, y0, y1)) or \
+                not all(math.isfinite(v) for v in h):
+            return {"fx": array.array("d", [0.0, 1.0]), "fy": array.array("d", [0.0, 1.0]),
+                    "flat": array.array("d", [0.0] * 4), "w": 2, "h": 2}
+        cut = self.get("cut")
+        x0, x1 = x0 - cut * math.sqrt(h[0]), x1 + cut * math.sqrt(h[0])
+        y0, y1 = y0 - cut * math.sqrt(h[2]), y1 + cut * math.sqrt(h[2])
+        flat = _backend.kde2d_grid(fx, fy, h, x0, x1, n, y0, y1, n)
+        gx = array.array("d", (x0 + (x1 - x0) * i / (n - 1) for i in range(n)))
+        gy = array.array("d", (y0 + (y1 - y0) * j / (n - 1) for j in range(n)))
+        return {"fx": gx, "fy": gy, "flat": flat, "w": n, "h": n}
+
+    def level_values(self) -> list[float]:
+        lv = self.get("levels")
+        if isinstance(lv, list):
+            return sorted(float(v) for v in lv)
+        g = self._grid()
+        values = sorted((v for v in g["flat"] if v == v), reverse=True)
+        total = sum(values)
+        if not total:
+            return [0.0]
+        thresh = self.get("thresh")
+        props = [thresh + (1.0 - thresh) * i / (lv - 1) for i in range(lv)] if lv > 1 \
+            else [thresh]
+        cum, out, acc, j = [], [], 0.0, 0
+        for v in values:
+            acc += v
+            cum.append(acc / total)
+        for p in props:                      # density whose upper set holds 1 − p
+            target = 1.0 - p
+            while j < len(cum) - 1 and cum[j] < target:
+                j += 1
+            out.append(values[min(j, len(values) - 1)])
+            j = 0
+        return sorted(set(out))
+
+    def _compute(self, sx: Any, sy: Any) -> dict[str, Any]:
+        g = self._grid()
+        lv = self.level_values()
+        lines: dict[float, Any] = {}
+        if self.get("lines"):
+            for level in lv:
+                lines[level] = stats.contour_lines(g["flat"], g["w"], g["h"], g["fx"], g["fy"],
+                                                   level)
+        bands = []
+        if self.get("filled"):
+            tops = lv[1:] + [math.inf]
+            for lo, hi in zip(lv, tops):
+                top = hi if math.isfinite(hi) else max(g["flat"]) * 1.0001 + 1e-300
+                rings = stats.contour_bands(g["flat"], g["w"], g["h"], g["fx"], g["fy"], lo, top)
+                if len(rings[0]):
+                    bands.append((lo, hi if math.isfinite(hi) else top, rings))
+        return {"levels": lv, "lines": lines, "bands": bands, "f_space": True}
+
+    def values(self, role: str) -> Any:
+        sx, sy = self._scales()
+        g = self._grid()
+        f, s = (g["fx"], sx) if role == "x" else (g["fy"], sy)
+        return array.array("d", [s.inverse(f[0]), s.inverse(f[-1])])
+
+    def _level_color(self, ctx: Any, i: int, level: float) -> Any:
+        if not self.get("filled") and not self.is_set("color_map") \
+                and self.get("colors") is None:
+            return self.base_color(ctx)
+        return super()._level_color(ctx, i, level)
+
+    def color_sources(self) -> list[tuple[Any, Any]]:
+        if not self.get("filled") and not self.is_set("color_map"):
+            return []
+        return super().color_sources()
+
+    def column_role(self, param: str) -> str | None:
+        return param if param in ("x", "y") else None
+
+    def legend_points(self) -> tuple[Any, Any] | None:
+        return self.px, self.py
+
+    def occupy(self, ctx: Any, occ: Any) -> None:
+        from . import Series
+        Series.occupy(self, ctx, occ)