plot-io 0.1.0__py3-none-any.whl

plot_io/scales.py

@@ -0,0 +1,93 @@
+"""Scales: a monotone function from data to a position line.
+
+A numeric scale maps a value ``v`` to ``f(v)``; an axis maps ``f(lo)…f(hi)`` linearly onto
+its length.  ``f`` is one of the functions pdf_io's ``arrays.scaled`` applies in C, so bulk
+data never goes through a Python loop.
+"""
+
+import math
+from collections.abc import Callable
+
+__all__ = ["Scale", "SCALES", "get_scale", "NUMERIC_SCALES", "AXIS_SCALES"]
+
+
+class Scale:
+    """``forward``/``inverse`` on single values; ``function`` is pdf_io's name for it."""
+
+    __slots__ = ("name", "function", "forward", "inverse", "lo", "hi", "options")
+
+    def __init__(self, name: str, function: str | None, forward: Callable[[float], float],
+                 inverse: Callable[[float], float], lo: float = -math.inf,
+                 hi: float = math.inf, options: dict | None = None) -> None:
+        self.name = name
+        self.function = function
+        self.forward = forward
+        self.inverse = inverse
+        self.lo = lo            # open domain bounds for log, closed for sqrt (see contains)
+        self.hi = hi
+        self.options = options or {}    # extra arguments of pdf_io's scaled (linthresh)
+
+    def contains(self, v: float) -> bool:
+        if not math.isfinite(v):
+            return False
+        if self.name in ("log", "log2", "log10"):
+            return v > 0
+        if self.name == "sqrt":
+            return v >= 0
+        if self.name == "exp":
+            return v < 709.78
+        if self.name == "logit":
+            return 0.0 < v < 1.0
+        return True
+
+    def f(self, v: float) -> float:
+        return self.forward(v) if self.contains(v) else math.nan
+
+    def __repr__(self) -> str:
+        return f"Scale({self.name})"
+
+
+def _safe_exp(t: float) -> float:
+    return math.exp(t) if t < 709.78 else math.inf
+
+
+SCALES = {
+    "linear": Scale("linear", None, lambda v: v, lambda t: t),
+    "log": Scale("log", "log", math.log, _safe_exp),
+    "log2": Scale("log2", "log2", math.log2, lambda t: 2.0 ** t if t < 1024 else math.inf),
+    "log10": Scale("log10", "log10", math.log10, lambda t: 10.0 ** t if t < 308 else math.inf),
+    "sqrt": Scale("sqrt", "sqrt", math.sqrt, lambda t: t * t),
+    "exp": Scale("exp", "exp", _safe_exp, lambda t: math.log(t) if t > 0 else -math.inf),
+}
+
+def _symlog(c: float) -> Scale:
+    def f(v: float) -> float:
+        return math.copysign(math.log10(1.0 + abs(v) / c), v)
+
+    def inv(t: float) -> float:
+        return math.copysign(c * (10.0 ** abs(t) - 1.0), t) if abs(t) < 308 else \
+            math.copysign(math.inf, t)
+    return Scale("symlog", "symlog", f, inv, options={"linthresh": c})
+
+
+def _logit_inverse(t: float) -> float:
+    if t > 300:
+        return 1.0
+    return 1.0 / (1.0 + 10.0 ** -t)
+
+
+SCALES["symlog"] = _symlog(1.0)
+SCALES["logit"] = Scale("logit", "logit", lambda p: math.log10(p / (1.0 - p)), _logit_inverse)
+
+NUMERIC_SCALES = tuple(SCALES)
+#: Scales an axis takes: the numeric ones, ``time`` (set by date data) and ``category``.
+AXIS_SCALES = NUMERIC_SCALES + ("time", "category")
+
+SCALES["time"] = Scale("time", None, lambda v: v, lambda t: t)
+
+
+def get_scale(name: str, linthresh: float | None = None) -> Scale:
+    """The scale ``name`` (``linthresh``: the linear region of ``symlog``)."""
+    if name == "symlog" and linthresh is not None and linthresh != 1.0:
+        return _symlog(float(linthresh))
+    return SCALES[name]

plot_io/series/__init__.py

@@ -0,0 +1,498 @@
+"""Series: data plus style, drawn in a plot.
+
+A series reports what the axes need (extent, categories, values for quantile limits),
+the color and size maps it uses, and draws itself and its legend glyph.  New plot types
+are new subclasses.
+"""
+
+import math
+from typing import Any
+
+from .. import _backend
+from ..settings import Setting, Settings, color, dpi_or_bool, number, route, simplify_value, table, text
+
+__all__ = ["Series", "COMMON", "DrawContext"]
+
+COMMON = dict(
+    name=lambda: Setting(text, doc="Legend entry; unset: not in the legend."),
+    alpha=lambda: Setting(number(0, 1), inherited=True, doc="Opacity, multiplying the colors' own."),
+    z=lambda: Setting(number(), doc="Drawing order (higher on top); default: call order."),
+    rasterize=lambda: Setting(dpi_or_bool, inherited=True,
+                              doc="`True` or a dpi: drawn as an image inside the PDF."),
+    simplify=lambda: Setting(simplify_value, inherited=True,
+                             doc="`\"auto\"`, `True`, a dpi or `False`: pdf_io's `simplify`."),
+    font_size=lambda: Setting(number(positive=True), inherited=True, doc="Base font size (pt).", length=True),
+    font_family=lambda: Setting(str, inherited=True, doc="Font family."),
+    text_color=lambda: Setting(color, inherited=True, doc="Text color."),
+    line_width=lambda: Setting(number(lo=0), inherited=True, doc="Line width (pt).", length=True),
+    color_map=lambda: Setting(_check_color_map, inherited=True,
+                              doc="How numbers in `color` become colors: a `ColorMap`, a map "
+                                  "name or colors (alias `cmap`); default: the style's."),
+    size_legend=lambda: Setting(_size_legend,
+                                doc="With a `size_map`: a size key in the legend (`True`), "
+                                    "none (`False`), or the values to show (a list)."),
+    size_map=lambda: Setting(_check_size_map,
+                             doc="How numbers in `size` become sizes: a `SizeMap` (alias "
+                                 "`smap`); without it, numbers are sizes in pt."),
+)
+
+
+def _size_legend(v: Any) -> Any:
+    if isinstance(v, bool):
+        return v
+    if isinstance(v, (list, tuple)) and v:
+        return [float(x) for x in v]
+    raise ValueError(f"expected True, False or a list of values, got {v!r}")
+
+
+def _check_color_map(v: Any) -> Any:
+    from ..values import check_color_map
+    return check_color_map(v)
+
+
+def _check_size_map(v: Any) -> Any:
+    from ..values import check_size_map
+    return check_size_map(v)
+
+
+#: Data arguments whose strings name columns when ``data`` is given (others than the
+#: positional ones).
+DATA_KEYS = ("color", "line_color", "marker_color", "size", "alpha", "group", "order",
+             "x_error", "y_error", "weights", "positions", "x", "y", "labels")
+#: Of those, the ones whose strings must be columns.
+MUST_BE_COLUMNS = ("group", "weights", "positions")
+#: Per-point keyword arguments split among the members of a group.
+SPLIT_KEYS = ("color", "line_color", "marker_color", "size", "alpha", "marker", "x_error",
+              "y_error", "order", "weights", "x", "y", "labels")
+
+
+def _positional(cls: type) -> list[str]:
+    import inspect
+    names = cls.__dict__.get("_POSITIONAL")
+    if names is None:
+        sig = inspect.signature(getattr(cls, "__init__"))
+        names = [n for n, p in sig.parameters.items()
+                 if n not in ("self", "plot") and p.kind in (p.POSITIONAL_ONLY,
+                                                             p.POSITIONAL_OR_KEYWORD)]
+        setattr(cls, "_POSITIONAL", names)
+    return names
+
+
+def _length(v: Any) -> int | None:
+    if v is None or isinstance(v, (str, bytes, dict, int, float)):
+        return None
+    try:
+        return len(v)
+    except TypeError:
+        return None
+
+
+class SeriesMeta(type):
+    """Resolves ``data`` columns and splits ``group`` calls before a series is made."""
+
+    def __call__(cls, *args: Any, **kw: Any) -> Any:
+        from .. import tables
+        from ..settings import canonical_kwargs
+        with_y = getattr(cls, "with_y", None)
+        if with_y is not None and (len(args) == 2 or "y" in kw):
+            return with_y()(*args, **kw)
+        names = _positional(cls)
+        kw = canonical_kwargs(kw, getattr(cls, "SETTINGS", {}), cls.__name__)
+        if len(args) > len(names):
+            raise TypeError(f"{cls.__name__}() takes {len(names)} positional arguments")
+        pos = dict(zip(names, args))
+        for n in names:
+            if n in kw:
+                if n in pos:
+                    raise TypeError(f"{cls.__name__}() got two values for {n!r}")
+                pos[n] = kw.pop(n)
+        plot = kw.get("plot")
+        data = kw.pop("data", None)
+        if data is None and plot is not None:
+            data = plot.get("data")
+        columns: dict[str, Any] = {}
+        orders: dict[str, Any] = {}
+        if data is not None:
+            for n, v in list(pos.items()):
+                if isinstance(v, str):
+                    pos[n], orders[n] = tables.column(data, v)
+                    columns[n] = v
+            for n in DATA_KEYS:
+                v = kw.get(n)
+                if isinstance(v, str) and n not in pos:
+                    if tables.has_column(data, v):
+                        if n in ("color", "line_color", "marker_color"):
+                            from ..settings import color
+                            try:
+                                color(v)
+                                tables.color_name_warning(v)
+                            except ValueError:
+                                pass
+                        kw[n], orders[n] = tables.column(data, v)
+                        columns[n] = v
+                    elif n in MUST_BE_COLUMNS:
+                        raise ValueError(f"no column {v!r} in data")
+        from .. import times as _times
+        time_params: dict[str, Any] = {}
+        for store in (pos, kw):
+            for n in ("x", "y", "positions"):
+                v = store.get(n)
+                if v is not None and _times.is_time_like(v):
+                    store[n], time_params[n] = _times.to_seconds(v)
+        if "group" in kw:
+            group = cls._grouped(pos, kw, columns, orders)
+            for m in group:
+                m.times = dict(time_params)
+            return group
+        for k in ("group_order", "group_missing", "group_colors"):
+            if k in kw:
+                raise ValueError(f"{k} applies with group=")
+        if getattr(cls, "LONG_FORM", False):
+            cls._long_form(pos, kw, columns, orders)
+        s = super().__call__(**pos, **kw)
+        s.columns = columns
+        s.times = time_params
+        return s
+
+    def _long_form(cls, pos: dict[str, Any], kw: dict[str, Any], columns: dict[str, Any],
+                   orders: dict[str, Any]) -> None:
+        """``values`` flat with categories per point in ``x`` (``y`` when horizontal):
+        a dict ``{category: values}`` in first-appearance (or categorical) order."""
+        key = "y" if kw.get("orientation") == "horizontal" else "x"
+        cats = kw.get(key)
+        vals = pos.get("values")
+        n = _length(vals)
+        if cats is None or n is None or isinstance(vals, dict) or _length(cats) != n:
+            return
+        kw.pop(key)
+        from .. import tables
+        out: dict[Any, list[Any]] = {}
+        for c in orders.get(key) or []:
+            out[c] = []
+        from .. import _backend
+        d = _backend.doubles(vals)
+        for c, v in zip(cats, d):
+            if tables.is_missing(c):
+                continue
+            out.setdefault(c, []).append(v)
+        pos["values"] = {k: v for k, v in out.items() if v}
+        if key in columns:
+            columns["categories"] = columns.pop(key)
+
+    def _grouped(cls, pos: dict[str, Any], kw: dict[str, Any], columns: dict[str, Any],
+                 orders: dict[str, Any]) -> Any:
+        from .. import tables
+        labels = kw.pop("group")
+        given = kw.pop("group_order", None)
+        missing = kw.pop("group_missing", "keep")
+        if missing not in ("keep", "drop"):
+            raise ValueError("group_missing is 'keep' or 'drop'")
+        colors = kw.pop("group_colors", None) or {}
+        n = _length(labels)
+        if n is None:
+            raise ValueError("group needs one label per point")
+        if hasattr(labels, "tolist"):
+            labels = labels.tolist()
+        keys, rows = tables.group_keys(labels, orders.get("group"), given, missing)
+        auto_marker = kw.get("marker") == "auto"
+        auto_style = kw.get("line_style") == "auto"
+        title = columns.get("group")
+        members: list[Any] = []
+        for k, (key, idx) in enumerate(zip(keys, rows)):
+            pk = {nm: (tables.take(v, idx) if _length(v) == n else v) for nm, v in pos.items()}
+            kk = {}
+            for nm, v in kw.items():
+                if nm in SPLIT_KEYS and _length(v) == n and not (
+                        isinstance(v, tuple) and len(v) in (3, 4)):
+                    v = tables.take(v, idx)
+                kk[nm] = v
+            if auto_marker:
+                kk["marker"] = tables.MARKER_CYCLE[k % len(tables.MARKER_CYCLE)]
+            if auto_style:
+                kk["line_style"] = tables.LINE_STYLE_CYCLE[k % len(tables.LINE_STYLE_CYCLE)]
+            if key in colors:
+                kk["color"] = colors[key]
+            kk.setdefault("name", str(key))
+            if k > 0 and kk.get("plot") is None and members:
+                kk["plot"] = members[0].plot
+            s = cls(**pk, **kk)
+            s.columns = dict(columns)
+            members.append(s)
+        group = tables.SeriesGroup(members, keys, title)
+        if members and title is not None:
+            members[0].plot._legend_title = title
+        return group
+
+
+#: Settings that may hold colors to map (numbers or a ``Density``).
+COLOR_SETTINGS = ("color", "line_color", "marker_color")
+#: Settings that may hold sizes to map.
+SIZE_SETTINGS = ("size",)
+
+
+def common(*names: str) -> dict[str, Setting]:
+    return {n: COMMON[n]() for n in names}
+
+
+class DrawContext:
+    """What a series needs while drawing: the canvas, the axes rectangle and mappers."""
+
+    def __init__(self, canvas: Any, plot: Any, rect: tuple[float, float, float, float],
+                 x: Any, y: Any, maps: dict[int, tuple[float, float]], palette: tuple) -> None:
+        self.canvas = canvas
+        self.plot = plot
+        self.rect = rect
+        self.x = x
+        self.y = y
+        x0, y0, x1, y1 = rect
+        self.xmap = x.mapper(x0, x1)
+        self.ymap = y.mapper(y1, y0)
+        self.maps = maps
+        self.palette = palette
+        self.figure_size = (x1, y1)
+
+    def color(self, value: Any) -> Any:
+        from ..style import resolve_color
+        return resolve_color(value, self.palette)
+
+    def map_limits(self, m: Any) -> tuple[float, float]:
+        return self.maps[id(m)]
+
+
+class Series(Settings, metaclass=SeriesMeta):
+    """Base of every series.  ``plot=None`` creates a plot (and its figure)."""
+
+    SETTINGS: dict[str, Setting] = table(**common("name", "alpha", "z", "rasterize",
+                                                  "simplify"))
+    #: Whether the series takes the next palette color when it has no color of its own.
+    USES_PALETTE = True
+    #: Whether the series is drawn on axes (``False``: pie charts, on plots without axes).
+    AXES = True
+    #: Dodge group of categorical series sharing category slots (``"box"``, ``"bar"``,
+    #: ``"points"``); ``None``: not dodged.
+    DODGE_GROUP: str | None = None
+    #: Stack group of series that stack with ``stack=True`` (``"bar"``, ``"area"``).
+    STACK_GROUP: str | None = None
+
+    #: Whether ``values`` may be flat with categories per point in ``x`` (long data).
+    LONG_FORM = False
+
+    def __init__(self, *, plot: Any = None, **settings: Any) -> None:
+        from ..figure import Figure
+        from ..plot import Plot
+
+        own, plot_kw, fig_kw = route(settings, [type(self), Plot, Figure],
+                                     extra={1: Plot.EXTRA})
+        if plot is None:
+            plot = Plot(**plot_kw, **fig_kw)
+        else:
+            if plot_kw:
+                plot.set(**plot_kw)
+            if fig_kw:
+                plot.figure.set(**fig_kw)
+        self.plot = plot
+        self._auto_color: Any = None
+        super().__init__(**own)
+        plot._add_series(self)
+
+    def _parent(self) -> Settings:
+        return self.plot
+
+    def set(self, **settings: Any) -> Any:
+        super().set(**settings)
+        self._cache: dict[Any, Any] = {}
+        return self
+
+    def cached(self, key: Any, compute: Any) -> Any:
+        """``compute()`` once until a setting changes (derived data: bins, densities…)."""
+        cache = self.__dict__.setdefault("_cache", {})
+        if key not in cache:
+            cache[key] = compute()
+        return cache[key]
+
+    def __repr__(self) -> str:
+        name = self.get("name")
+        return f"{self.kind}({name!r})" if name else f"{self.kind}()"
+
+    # -- what the axes need --------------------------------------------------------
+
+    def has_axis(self, role: str) -> bool:
+        return True
+
+    def time_role(self, role: str) -> Any:
+        """``(True, tz)`` when the data on axis ``role`` were dates, else ``None``."""
+        for param, tz in getattr(self, "times", {}).items():
+            if self.column_role(param) == role:
+                return (True, tz)
+        return None
+
+    def column_role(self, param: str) -> str | None:
+        """The axis a data argument named ``param`` is drawn on (for axis titles taken from
+        column names)."""
+        if param in ("x", "y"):
+            return param
+        if param in ("values", "height", "categories"):
+            cat = getattr(self, "_cat_role", None)
+            if cat is not None:
+                role = cat()
+                return role if param == "categories" else ("y" if role == "x" else "x")
+            horizontal = "orientation" in self.SETTINGS and \
+                self.get("orientation") == "horizontal"
+            return "y" if horizontal else "x"
+        return None
+
+    def axis_title(self, role: str) -> Any:
+        """The column name drawn on axis ``role``, if the series was given a column."""
+        for param, name in getattr(self, "columns", {}).items():
+            if self.column_role(param) == role:
+                return name
+        return None
+
+    def categories(self, role: str) -> list[Any] | None:
+        """Categories on axis ``role``, or ``None`` for numeric data."""
+        return None
+
+    def values(self, role: str) -> Any:
+        """All data values on axis ``role`` (for quantile limits)."""
+        return None
+
+    def extent(self, role: str, scale: Any) -> tuple[int, float, float]:
+        v = self.values(role)
+        if v is None:
+            return 0, math.nan, math.nan
+        return _backend.extent(v, scale.function)
+
+    def inverts(self, role: str) -> bool:
+        """Whether axis ``role`` is inverted by default for this series."""
+        return False
+
+    def zero_default(self, role: str) -> bool:
+        """Whether axis ``role`` includes 0 under ``include_zero="auto"`` (Vega-Lite)."""
+        return False
+
+    def zero_padding(self, role: str) -> bool:
+        """Whether automatic limits on axis ``role`` take no padding."""
+        return False
+
+    def color_sources(self) -> list[tuple[Any, Any]]:
+        """``(ColorMap, values)`` pairs the series maps: its color settings holding numbers
+        or a ``Density``, through ``color_map_obj()``."""
+        from ..values import Values, is_computed
+        out = []
+        for name in COLOR_SETTINGS:
+            if name in self.SETTINGS:
+                v = self.get(name)
+                if isinstance(v, Values) or is_computed(v):
+                    out.append((self.color_map_obj(), self.point_values(v)))
+        return out
+
+    def size_sources(self) -> list[tuple[Any, Any]]:
+        out = []
+        for name in SIZE_SETTINGS:
+            m = self.size_map_obj(name)
+            if m is not None:
+                out.append((m, self.point_values(self.get(name))))
+        return out
+
+    def color_map_obj(self) -> Any:
+        """The ``ColorMap`` numbers in ``color`` go through: ``color_map`` when it is a
+        ``ColorMap`` (shared), else one made from it for this series."""
+        from ..maps import ColorMap
+        m = self.get("color_map") if "color_map" in self.SETTINGS else self.plot.get("color_map")
+        if isinstance(m, ColorMap):
+            return m
+        return self.cached(("color_map", repr(m)), lambda: ColorMap(m))
+
+    def size_map_obj(self, name: str = "size") -> Any:
+        """The ``SizeMap`` of a size setting holding values, or ``None`` (sizes in pt)."""
+        from ..maps import SizeMap
+        from ..values import is_computed
+        if name not in self.SETTINGS or "size_map" not in self.SETTINGS:
+            return None
+        v = self.get(name)
+        m = self.get("size_map")
+        if is_computed(v):
+            return m if m is not None else self.cached("size_map", SizeMap)
+        if m is not None and not isinstance(v, (int, float)):
+            return m
+        return None
+
+    def point_values(self, v: Any) -> Any:
+        """The numbers of a value setting: given ``Values``, sizes, or a computed
+        ``Density``."""
+        from ..values import Values, is_computed
+        if isinstance(v, Values):
+            return v.data
+        if is_computed(v):
+            return self.cached(("density", id(v)), lambda: v.compute(self))
+        return v
+
+    def palette_count(self) -> int:
+        """Palette colors the series takes when it has no color (one per value set of
+        multi-set series)."""
+        return 1
+
+    def auto_colors(self, ctx: Any) -> list[Any]:
+        """The palette colors given to the series (``palette_count()`` of them)."""
+        cols = getattr(self, "_auto_colors", None) or [self._auto_color]
+        return [ctx.color(c) for c in cols]
+
+    def has_own_color(self) -> bool:
+        return self.is_set("color")
+
+    def base_color(self, ctx: DrawContext) -> Any:
+        """The series color: its own, or the palette color it was given."""
+        c = self.get("color") if "color" in self.SETTINGS else None
+        if c is None:
+            c = self._auto_color
+        return ctx.color(c)
+
+    # -- drawing ---------------------------------------------------------------------
+
+    def draw(self, ctx: DrawContext) -> None:
+        raise NotImplementedError
+
+    def draw_overlay(self, ctx: DrawContext) -> None:
+        """Drawn after every series, outside the axes clip (value labels)."""
+
+    def legend_glyph(self, ctx: DrawContext, x: float, y: float, w: float, h: float) -> None:
+        """Draw the legend symbol in the box ``(x, y, w, h)`` (y: top)."""
+
+    def legend_colors(self, ctx: DrawContext) -> tuple[Any, Any]:
+        """``(fill, stroke)`` of simple legend symbols (``Legend.glyph``)."""
+        return self.base_color(ctx), None
+
+    def legend_labels(self) -> list[Any]:
+        """Legend entries of the object: its ``name``, or none."""
+        name = self.get("name")
+        return [] if name is None else [name]
+
+    def legend_glyph_at(self, ctx: Any, i: int, x: float, y: float, w: float, h: float) -> None:
+        """The symbol of legend entry ``i`` (series with several entries override this)."""
+        self.legend_glyph(ctx, x, y, w, h)
+
+    def legend_colors_at(self, ctx: Any, i: int) -> tuple[Any, Any]:
+        return self.legend_colors(ctx)
+
+    def occupy(self, ctx: DrawContext, occ: Any) -> None:
+        """Marks what the series draws on ``occ`` (an :class:`~plot_io.occupancy.Occupancy`),
+        for legend placement; by default its legend points."""
+        pts = self.legend_points()
+        if pts is not None:
+            from ..occupancy import thin
+            xs, ys = thin(*(_backend.doubles(v) for v in pts))
+            size = self.get("size") if "size" in self.SETTINGS else None
+            radius = size / 2 if isinstance(size, (int, float)) else 0.0
+            occ.points(ctx.xmap.many(xs), ctx.ymap.many(ys), radius)
+
+    def occupy_all(self, ctx: DrawContext, occ: Any) -> None:
+        """Marks the whole axes (series that fill them: heatmaps, contours)."""
+        occ.rects([ctx.rect[0]], [ctx.rect[1]], [ctx.rect[2]], [ctx.rect[3]])
+
+    def legend_points(self) -> tuple[Any, Any] | None:
+        """``(xs, ys)`` data points for the legend's ``auto`` placement."""
+        return None
+
+    def simplify_arg(self, lines: bool = False) -> Any:
+        """pdf_io's ``simplify`` for this series' calls (``lines``: lines and areas)."""
+        return _backend.simplify_arg(self.get("simplify"), lines)