plot-io 0.12.0__py3-none-win_arm64.whl

plot_io/density.py

@@ -0,0 +1,113 @@
+"""Point density: a 2-D Gaussian kernel density estimate at each point of a series, as
+scipy's ``gaussian_kde(np.array([x, y]))(np.array([x, y]))`` computes it.
+
+The kernel covariance is the data covariance scaled by Scott's factor n^(-1/6) (or
+Silverman's), and the density is computed on the axis positions (log axes in log space).
+plot_io's C library does the work: an exact sum for small inputs, a binned and blurred grid for
+large ones.
+"""
+
+import math
+from typing import Any
+
+from vec_io.arrays import doubles
+
+from . import kernels
+from .settings import number, one_of
+from .validators import real
+
+__all__ = ["Density"]
+
+_check_normalize = one_of("pdf", "count", "max")
+_check_method = one_of("auto", "exact", "grid")
+
+
+def _bandwidth(v: Any) -> Any:
+    if isinstance(v, str):
+        return one_of("scott", "silverman")(v)
+    if real(v) is None and hasattr(v, "__iter__"):
+        m = [list(map(float, row)) for row in v]
+        if len(m) != 2 or any(len(r) != 2 for r in m):
+            raise ValueError("a bandwidth matrix is 2 × 2")
+        if abs(m[0][1] - m[1][0]) > 1e-12 * max(1.0, abs(m[0][1])):
+            raise ValueError("a bandwidth matrix is symmetric")
+        if not (m[0][0] > 0 and m[1][1] > 0 and m[0][0] * m[1][1] - m[0][1] ** 2 > 0):
+            raise ValueError("a bandwidth matrix is positive definite")
+        return (m[0][0], m[0][1], m[1][1])
+    return number(positive=True)(v)
+
+
+class Density:
+    """The local density of a series' points, as values for ``color``, ``size``, ``opacity``
+    or ``draw_order``.
+
+    ``bandwidth``: ``"scott"`` (default), ``"silverman"``, a number (isotropic standard
+    deviation in axis units) or a 2 × 2 covariance matrix; ``adjust`` multiplies it.
+    ``normalize``: ``"pdf"`` (scipy's values), ``"count"`` (points per unit area) or
+    ``"max"`` (0–1).  ``method``: ``"exact"``, ``"grid"`` or ``"auto"`` (exact up to
+    20,000 points); ``grid``: grid size of the grid method.
+
+    As scipy's ``gaussian_kde(np.array([x, y]))(np.array([x, y]))``: the kernel covariance
+    is the data covariance times Scott's factor n^(-1/6), on the axis positions (log axes in
+    log space); NaN points get NaN.  ``"exact"`` sums over all points (multithreaded, equal
+    to scipy's); ``"grid"`` bins, blurs and reads back (about 1e-4 relative).  Accepted by
+    ``color``, ``edge_color``, ``size`` (through the size axis), ``opacity`` (mapped to
+    ``opacity_range``), ``draw_order`` (``"density"``: densest last), and the marker colors
+    of Line, Errorbar, Stem and Regression; computed per series, cached until a setting
+    changes, in ``series.stats["density"]``.
+    """
+
+    def __init__(self, bandwidth: Any = "scott", adjust: float = 1.0, normalize: str = "pdf",
+                 method: str = "auto", grid: int = 512) -> None:
+        self.bandwidth = _bandwidth(bandwidth)
+        self.adjust = number(positive=True)(adjust)
+        self.normalize = _check_normalize(normalize)
+        self.method = _check_method(method)
+        self.grid = int(number(lo=16, hi=4096, integer=True)(grid))
+
+    def __repr__(self) -> str:
+        return f"Density({self.bandwidth!r}, normalize={self.normalize!r})"
+
+    def label(self) -> str:
+        return {"pdf": "density", "count": "points per unit area",
+                "max": "relative density"}[self.normalize]
+
+    def compute(self, series: Any) -> Any:
+        """The density at each point of ``series`` (its ``density_points()``)."""
+        xs, ys = series.density_points()
+        return self.at(xs, ys)
+
+    def at(self, xs: Any, ys: Any) -> Any:
+        """The density at each point ``(xs[i], ys[i])`` (positions, not data)."""
+        xs, ys = doubles(xs), doubles(ys)
+        h = self.covariance(xs, ys)
+        if not all(math.isfinite(v) for v in h):          # fewer than 2 distinct points
+            import array
+            import warnings
+            warnings.warn("density of fewer than 2 distinct points: drawn in one color",
+                          stacklevel=2)
+            return array.array("d", (1.0 if math.isfinite(x) and math.isfinite(y)
+                                     else math.nan for x, y in zip(xs, ys)))
+        dens = kernels.kde2d_at(xs, ys, h, self.method, self.grid)
+        return self._normalized(dens, xs, ys)
+
+    def covariance(self, xs: Any, ys: Any, weights: Any = None) -> tuple[float, float, float]:
+        """The kernel covariance ``(h11, h12, h22)`` (``weights``: as repeated points)."""
+        a = self.adjust * self.adjust
+        if isinstance(self.bandwidth, tuple):
+            h11, h12, h22 = self.bandwidth
+            return (h11 * a, h12 * a, h22 * a)
+        if isinstance(self.bandwidth, float):
+            s = self.bandwidth * self.bandwidth * a
+            return (s, 0.0, s)
+        h11, h12, h22 = kernels.kde2d_bandwidth(xs, ys, self.bandwidth, weights)
+        return (h11 * a, h12 * a, h22 * a)
+
+    def _normalized(self, dens: Any, xs: Any, ys: Any) -> Any:
+        if self.normalize == "pdf":
+            return dens
+        if self.normalize == "count":
+            n = sum(1 for x, y in zip(xs, ys) if x == x and y == y)
+            return kernels.scaled(dens, float(n))
+        _, _, hi = kernels.extent(dens, None)
+        return kernels.scaled(dens, 1.0 / hi if hi and math.isfinite(hi) and hi > 0 else 1.0)

plot_io/ext.py

@@ -0,0 +1,68 @@
+"""Plot types of your own.  Subject to change while plot_io is 0.x.
+
+A type is a :class:`Series` subclass: its data arguments (``DATA``: argument → slot kind,
+:data:`KINDS`), its settings (``SETTINGS``: :func:`table` of :class:`Setting`, with
+``default=``; :func:`common` for the shared ones) and its declaration (``SPEC``:
+:class:`Spec` of :class:`Lay` layers drawing :class:`Table` columns with a geom of
+:data:`GEOMS` and a paint role).  :func:`register` adds ``plot.<name>(…)``.  Paint roles and
+layer options are checked at the type's first drawing.  A layer's ``stat`` (a name of
+:data:`STATS`; :func:`stat` registers one, with ``defaults`` for the settings it reads that a
+type does not have) runs once per call: table functions ``fn(series, build)`` read its output
+in ``build.stats[name]``.  A layer's ``position`` (``"stack"``, ``"dodge"``…) is asked of its
+tables holding ``keys`` and ``heights`` (``bottom`` and ``top`` written).  Unknown stats and
+positions raise.  :func:`auto_colors` gives the
+colors of the plot's a series took (``call.color_count`` of them, set in a spec's ``finish``);
+:func:`base_colors` each set's color (``color`` given as :func:`set_colors` takes it, else
+the plot's colors); ``register(…, levels="split")`` lets ``color=``, ``marker=``… take a
+column of levels (one call per level, as ``group=``).
+:func:`norm_quantile` and :func:`beta_quantile` are quantile functions for stats;
+:func:`axis_of` and :func:`axis_scale` give the axis a data role is drawn on and its scale
+(stats computed in the axis' space); :func:`group` embeds a setting group (``"Marker"``);
+:func:`with_alpha` scales a color's opacity; :class:`PanelSet` (``plot_io.facets``) calls a
+type on every panel of a figure with the panel's arguments.  ``docs/api.md``,
+"Plot types of your own", has an example.
+"""
+
+from typing import Any
+
+from .color.values import color, with_alpha
+from .columns import KINDS, Columns
+from .facets import PanelSet
+from .layers.base import Series, common
+from .layers.core import Table
+from .layers.geoms import GEOMS, LAYER_OPTIONS
+from .layers.grammar import Lay, Spec, axis_of
+from .layers.paint import ROLES, Paint, auto_colors, base_colors, role
+from .layers.stats import STATS, axis_scale, beta_quantile, compute, norm_quantile, stat
+from .presets._util import set_color_value as set_colors
+from .settings import Setting, table
+from .validators import angle, boolean, either, length, literal, number, one_of, pair, text
+
+__all__ = ["Series", "Spec", "Lay", "Table", "Columns", "KINDS", "GEOMS", "LAYER_OPTIONS",
+           "ROLES", "Paint", "role", "auto_colors", "base_colors", "set_colors", "STATS",
+           "stat", "compute",
+           "norm_quantile", "beta_quantile", "axis_of", "axis_scale", "group", "with_alpha",
+           "PanelSet", "Setting", "table", "common",
+           "color", "number", "length", "angle", "one_of", "literal", "either", "pair",
+           "boolean", "text", "register"]
+
+
+def register(name: str, cls: Any, levels: str | None = None) -> None:
+    """Add ``plot.<name>(…)``, ``figure.<name>(…)`` and ``facets.<name>(…)`` drawing
+    ``cls`` (``plt.registry.layers.add``); ``levels``: ``"split"`` (one call per level
+    combination, as ``group=`` members) or ``"items"`` (each item by its levels) makes the
+    type's ``color``, ``marker``, ``line_style`` and ``size`` take levels."""
+    from .levels import adopt
+    from .registry import layers
+    if levels is not None:
+        if levels not in ("split", "items"):
+            raise ValueError(f"levels: expected 'split' or 'items', got {levels!r}")
+        adopt(cls, levels)
+    layers.add(name, cls)
+
+
+def group(name: str, prefix: str = "", **options: Any) -> dict[str, Setting]:
+    """The settings of group ``name`` (``"Marker"``, ``"Stroke"``, ``"FillStyle"``,
+    ``"TextStyle"``, ``"Font"``) with ``prefix`` (``groups.<name>.embed``)."""
+    from . import groups
+    return getattr(groups, name).embed(prefix, **options)

plot_io/facets.py

@@ -0,0 +1,240 @@
+"""Facets: small multiples of one table, split by the values of one or two columns.
+
+``fig, facets = plt.facet(data, col=…, row=…)`` makes an ordinary :class:`~plot_io.Figure`
+with one plot per combination of values, and a :class:`FacetSet` of those panels.  A plot
+type called on the set (``facets.scatter("x", "y")``) adds one layer per panel, on that
+panel's rows, and returns a :class:`LayerSet`; group levels keep their colors across panels.
+Axes are shared through the figure's ``share``, color and size axes through its
+``share_color`` and ``share_size`` (``"all"`` by default), as on any figure; the panels'
+legend entries go into one figure legend (``share_legend="all"`` by default).
+"""
+
+from typing import Any, Callable
+
+from ._methods import PanelSetMixin
+from .figure import Figure
+from .tables import column, group_keys, take
+
+__all__ = ["facet", "PanelSet", "FacetSet", "LayerSet"]
+
+#: Decorations a facet set adds to every panel.
+_DECORATIONS = ("hline", "vline", "axline", "hspan", "vspan", "text", "annotate",
+                "bracket", "brackets")
+_SIDES = ("right", "left", "top", "bottom")
+
+
+def subset(data: Any, idx: list[int]) -> Any:
+    """The rows ``idx`` of a table."""
+    if hasattr(data, "iloc"):                        # pandas
+        return data.iloc[idx]
+    if isinstance(data, dict):
+        return {k: take(v, idx) for k, v in data.items()}
+    if getattr(getattr(data, "dtype", None), "names", None) is not None:
+        return data[idx]                             # numpy structured array
+    try:
+        return data[idx]                             # polars and others
+    except (TypeError, KeyError, IndexError):
+        raise ValueError("facets need a dict, a data frame or a structured array") from None
+
+
+def _levels(data: Any, name: str | None, order: Any) -> tuple[list[Any], list[list[int]]]:
+    if name is None:
+        n = len(next(iter(data.values()))) if isinstance(data, dict) else len(data)
+        return [None], [list(range(n))]
+    values, cat_order = column(data, name)
+    if hasattr(values, "tolist"):
+        values = values.tolist()
+    given = order.get(name) if isinstance(order, dict) else None
+    return group_keys(values, cat_order, given, "drop")
+
+
+class LayerSet(list):
+    """The layers (or decorations) one facet-set call added, one per panel."""
+
+    def set(self, **settings: Any) -> "LayerSet":
+        """Set settings on every member."""
+        for s in self:
+            s.set(**settings)
+        return self
+
+    @property
+    def plots(self) -> list[Any]:
+        return [getattr(s, "plot", None) for s in self]
+
+
+class PanelSet(PanelSetMixin):
+    """Panels of a figure: plot types and decorations called on the set are added to every
+    panel, with the panel's arguments (:meth:`arguments`).  Base of :class:`FacetSet`
+    (each panel's rows of one table) and of the panel sets of ``plt.pairs`` (each panel's
+    variables)."""
+
+    def __init__(self, figure: Figure, data: Any, panels: dict[Any, Any],
+                 colorbar: Any = None) -> None:
+        self.figure = figure
+        self.data = data
+        #: panel key → plot
+        self.panels = panels
+        self._colorbar_side = colorbar
+        self._group_levels: dict[str, list[Any]] = {}
+        #: the colorbar of the first call mapping colors (``colorbar=`` given)
+        self.colorbar: Any = None
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({len(self.panels)} panels)"
+
+    def __iter__(self) -> Any:
+        return iter(self.panels.values())
+
+    def __len__(self) -> int:
+        return len(self.panels)
+
+    def plot(self, key: Any) -> Any:
+        """The panel of key ``key``."""
+        return self.panels[key]
+
+    def part(self, key: Any) -> Any:
+        """What panel ``key`` draws (its key here; a facet's rows, a pair's variables)."""
+        return key
+
+    def arguments(self, key: Any, args: tuple, kw: dict[str, Any]) -> tuple[tuple, dict]:
+        """The arguments of a call on panel ``key``: the call's, with ``data=`` the set's
+        table."""
+        if self.data is not None:
+            kw = {**kw, "data": self.data}
+        return args, kw
+
+    def map(self, func: Callable[[Any, Any], Any]) -> list[Any]:
+        """``func(plot, part)`` for every panel (:meth:`part`); returns the results."""
+        return [func(p, self.part(k)) for k, p in self.panels.items()]
+
+    def call(self, method: str, args: tuple, kw: dict[str, Any]) -> LayerSet:
+        """``plot.<method>(*args, **kw)`` on every panel, with the panel's arguments."""
+        kw = dict(kw)
+        if self.data is not None:
+            from .levels import facet_orders
+            facet_orders(self.data, kw)
+        group = kw.get("group")
+        if isinstance(group, str) and self.data is not None:
+            levels = self._group_levels.get(group)
+            if levels is None:
+                values, cat = column(self.data, group)
+                if hasattr(values, "tolist"):
+                    values = values.tolist()
+                levels = group_keys(values, cat, kw.get("group_order"),
+                                    kw.get("group_missing", "keep"))[0]
+                self._group_levels[group] = levels
+            kw.setdefault("group_order", levels)
+            colors = {lv: f"C{i}" for i, lv in enumerate(levels)}
+            colors.update(kw.get("group_colors") or {})
+            kw["group_colors"] = colors
+        out = LayerSet()
+        for k, p in self.panels.items():
+            a, w = self.arguments(k, args, kw)
+            out.append(getattr(p, method)(*a, **w))
+        if self._colorbar_side and self.colorbar is None and out:
+            self._add_colorbar(out[0])
+        return out
+
+    def _add_colorbar(self, first: Any) -> None:
+        """The colorbar of the panels' color axes, in the lanes beside them, for the first
+        call that maps colors."""
+        from .layers.paint import color_sources
+        members = list(first) if isinstance(first, list) else [first]
+        mapped = next((s for s in members if hasattr(s, "SETTINGS") and color_sources(s)),
+                      None)
+        if mapped is None:
+            return
+        self.colorbar = self.figure.add_colorbar(mapped, side=self._colorbar_side)
+
+
+class FacetSet(PanelSet):
+    """The panels of a facet figure (``plt.facet``), keyed by facet value (or ``(row
+    value, col value)``): plot types and decorations called on it are added to every panel
+    (on its rows); :meth:`map` passes each panel's rows."""
+
+    def __init__(self, figure: Figure, data: Any, panels: dict[Any, Any],
+                 rows: dict[Any, list[int]], colorbar: Any) -> None:
+        super().__init__(figure, data, panels, colorbar)
+        self._rows = rows
+
+    def part(self, key: Any) -> Any:
+        """The rows of panel ``key``."""
+        return subset(self.data, self._rows[key])
+
+    def arguments(self, key: Any, args: tuple, kw: dict[str, Any]) -> tuple[tuple, dict]:
+        return args, {**kw, "data": self.part(key)}
+
+    def subsets(self) -> list[tuple[Any, Any, Any]]:
+        """``(key, plot, rows)`` of every panel."""
+        return [(k, p, self.part(k)) for k, p in self.panels.items()]
+
+
+def _decoration(name: str) -> Any:
+    def method(self: PanelSet, *args: Any, **settings: Any) -> LayerSet:
+        return LayerSet(getattr(p, name)(*args, **settings) for p in self.panels.values())
+    method.__name__ = name
+    method.__qualname__ = f"PanelSet.{name}"
+    method.__doc__ = f"``plot.{name}(…)`` on every panel."
+    return method
+
+
+for _name in _DECORATIONS:
+    setattr(PanelSet, _name, _decoration(_name))
+
+
+def facet(data: Any, col: str | None = None, row: str | None = None, wrap: int | None = None,
+          share: Any = "grid", order: Any = None, title_format: str = "{name} = {value}",
+          colorbar: Any = None, **figure_settings: Any) -> tuple[Figure, FacetSet]:
+    """Small multiples of ``data``: a figure with one plot per value of ``col`` (and
+    ``row``), and the :class:`FacetSet` of its panels.  ``wrap`` wraps ``col`` into rows of
+    that many panels; ``order``: ``{column: [values]}``; ``colorbar``: ``"right"``,
+    ``"left"``, ``"top"`` or ``"bottom"``, one colorbar beside the panels (in the lanes,
+    ``figure.add_colorbar``) for their color axes (shared: ``share_color="all"`` unless
+    given; size axes likewise, ``share_size="all"``, with one size key; one legend for the
+    panels' entries on the right: ``share_legend="all"`` unless given).
+
+    Example::
+
+        rng = random.Random(1)
+        df = {"x": [rng.gauss(0, 1) for _ in range(150)], "condition": ["a", "b", "c"] * 50}
+        df["y"] = [v + rng.gauss(0, 0.5) for v in df["x"]]
+        fig, panels = plt.facet(df, col="condition", width=120, height=120)
+        panels.scatter("x", "y")
+    """
+    if col is None and row is None:
+        raise ValueError("facet needs col, row or both")
+    if wrap is not None and (row is not None or col is None):
+        raise ValueError("wrap applies to col alone")
+    if colorbar not in (None, False) and colorbar not in _SIDES:
+        raise ValueError(f"colorbar is None or one of {', '.join(map(repr, _SIDES))}")
+    cols, col_rows = _levels(data, col, order)
+    rows, row_rows = _levels(data, row, order)
+    panels: list[tuple[Any, Any, list[int]]] = []
+    for rk, ri in zip(rows, row_rows):
+        rset = set(ri)
+        for ck, ci in zip(cols, col_rows):
+            panels.append((rk, ck, [i for i in ci if i in rset]))
+    if wrap:
+        ncol = wrap
+        nrow = -(-len(panels) // ncol)
+    else:
+        nrow, ncol = len(rows), len(cols)
+    grid: list[list[Any]] = [[None] * ncol for _ in range(nrow)]
+    ids = []
+    for k in range(len(panels)):
+        r, c = divmod(k, ncol)
+        grid[r][c] = f"panel{k}"
+        ids.append(f"panel{k}")
+    figure_settings.setdefault("share_color", "all")
+    figure_settings.setdefault("share_size", "all")
+    figure_settings.setdefault("share_legend", "all")
+    fig = Figure(layout=grid, share=share, **figure_settings)
+    by_key: dict[Any, Any] = {}
+    idx_of: dict[Any, list[int]] = {}
+    for pid, (rk, ck, idx) in zip(ids, panels):
+        key = ck if row is None else (rk if col is None else (rk, ck))
+        parts = [(n, v) for n, v in ((row, rk), (col, ck)) if n is not None]
+        title = ", ".join(title_format.format(name=n, value=v) for n, v in parts)
+        by_key[key] = fig.plot(pid, title=title)
+        idx_of[key] = idx
+    return fig, FacetSet(fig, data, by_key, idx_of, colorbar or None)