plot-io 0.12.0__py3-none-win_amd64.whl

plot_io/columns.py

@@ -0,0 +1,614 @@
+"""Data intake: the data arguments of a plot type's call as columns.
+
+One rule for every plot type:
+
+* a column is a sequence, a 1-D buffer, or a pandas or polars series (nullable arrays
+  included: their missing values are NaN); ``None`` and NaN are missing values; a
+  generator is read once;
+* numeric slots (values, heights, ``y``, weights) take numbers: text is refused, naming
+  the argument (``"1"`` is not read as 1); durations (``timedelta``, ``timedelta64``)
+  are read as seconds;
+* position slots (``x``, ``y`` of the (x, y) types) take numbers or categories: text,
+  or a pandas or polars categorical (its order kept), gives :class:`Keys`;
+* several columns: a dict ``{name: column}``, a data frame (columns by name), a 2-D
+  buffer (its **columns** are the variables, as in numpy and pandas), or a Python list or
+  tuple of sequences (each a column).
+
+:class:`Columns` reads the data arguments of one call once, from the type's ``DATA``
+declaration (each argument with its slot kind), into columns and groups of rows.
+"""
+
+import array
+import datetime as _dt
+import math
+import numbers as _numbers
+from typing import Any
+
+from vec_io.arrays import buffer_shape, doubles, grid
+
+from .validators import typed
+
+__all__ = ["Keys", "Columns", "Group", "KINDS", "numbers", "positions", "is_categorical",
+           "is_frame", "is_several", "sets", "named_sets", "weights"]
+
+_NAN = math.nan
+
+
+def is_frame(v: Any) -> bool:
+    """Whether ``v`` is a data frame (pandas, polars: named columns)."""
+    shape = getattr(v, "shape", None)
+    return (hasattr(v, "columns") and not isinstance(v, dict) and shape is not None
+            and len(shape) == 2 and hasattr(v, "__getitem__"))
+
+
+def _shape(v: Any) -> tuple[int, ...] | None:
+    """The shape of a numeric buffer, ``None`` for other values (durations included)."""
+    if isinstance(v, (str, bytes, list, tuple, dict)) \
+            or getattr(getattr(v, "dtype", None), "kind", None) == "m":
+        return None
+    return buffer_shape(v)
+
+
+def _missing(item: Any) -> bool:
+    return item is None or type(item).__name__ in ("NAType", "NaTType", "NoneType")
+
+
+def _plain(v: Any) -> Any:
+    """A pandas or polars series (or pandas array) as a numpy array, missing values of a
+    nullable numeric type as NaN, durations as seconds; other values unchanged."""
+    if getattr(getattr(v, "dtype", None), "kind", None) == "m":
+        return _durations(v)
+    if not hasattr(v, "to_numpy") or hasattr(v, "__array_interface__"):
+        return v
+    try:
+        return v.to_numpy(dtype=float, na_value=_NAN)
+    except (TypeError, ValueError):
+        return v.to_numpy()
+
+
+def _durations(v: Any) -> array.array:
+    """A ``timedelta64`` array or series as seconds (``NaT``: NaN)."""
+    if hasattr(v, "to_numpy") and not hasattr(v, "__array_interface__"):
+        v = v.to_numpy()
+    ns = v.astype("timedelta64[ns]").astype("int64").tolist()
+    return array.array("d", (_NAN if n == -(2 ** 63) else n / 1e9 for n in ns))
+
+
+def numbers(v: Any, what: str) -> array.array:
+    """One numeric column as ``array('d')`` (missing values NaN); text refused."""
+    if isinstance(v, array.array) and v.typecode == "d":
+        return v
+    if isinstance(v, (str, bytes)):
+        raise ValueError(f"{what} must be numbers, not text {v!r}")
+    if isinstance(v, Keys):
+        raise ValueError(f"{what} must be numbers, not categories")
+    v = _plain(v)
+    shape = _shape(v)
+    if shape is not None:
+        if len(shape) != 1:
+            raise ValueError(f"{what} must be 1-D, not of shape {shape}")
+        return doubles(v)
+    if hasattr(v, "tolist") and not isinstance(v, (list, tuple)):
+        v = v.tolist()
+    if isinstance(v, (list, tuple)):
+        try:
+            return array.array("d", v)              # numbers only (no text, no None)
+        except TypeError:
+            pass
+    out = array.array("d")
+    append = out.append
+    for item in v:
+        if isinstance(item, float):
+            append(item)
+        elif isinstance(item, _numbers.Real) and not isinstance(item, bool):
+            append(float(item))
+        elif _missing(item):
+            append(_NAN)
+        elif isinstance(item, _dt.timedelta):
+            append(item.total_seconds())
+        elif isinstance(item, (str, bytes)):
+            raise ValueError(f"{what} must be numbers, not text: {item!r}")
+        else:
+            try:
+                append(float(item))
+            except (TypeError, ValueError):
+                raise ValueError(f"{what} must be numbers, not {type(item).__name__}: "
+                                 f"{item!r}") from None
+    return out
+
+
+class Keys:
+    """Positions given as categories: the key of each point (``None``: missing) and the
+    categories in order (a categorical's order, else first appearance)."""
+
+    __slots__ = ("keys", "levels", "own")
+
+    def __init__(self, keys: Any, order: Any = None) -> None:
+        self.keys = [None if _missing(k) or (isinstance(k, float) and math.isnan(k)) else k
+                     for k in keys]
+        levels = list(order) if order is not None else []
+        seen = set(levels)
+        for k in self.keys:
+            if k is not None and k not in seen:
+                seen.add(k)
+                levels.append(k)
+        self.levels = levels
+        index = {k: float(i) for i, k in enumerate(levels)}
+        #: each point at its category's index among ``levels`` (before the axes are known)
+        self.own = array.array("d", (_NAN if k is None else index[k] for k in self.keys))
+
+    def __len__(self) -> int:
+        return len(self.keys)
+
+    def __iter__(self) -> Any:
+        return iter(self.keys)
+
+    def __getitem__(self, i: Any) -> Any:
+        return self.keys[i]
+
+    def take(self, idx: list[int]) -> "Keys":
+        out = Keys.__new__(Keys)
+        out.keys = [self.keys[i] for i in idx]
+        out.levels = self.levels
+        out.own = array.array("d", (self.own[i] for i in idx))
+        return out
+
+    def on(self, categories: list[Any]) -> array.array:
+        """Each point at its category's index on an axis with ``categories``."""
+        index: dict[Any, float] = {}
+        for i, c in enumerate(categories):
+            index.setdefault(c, float(i))
+        return array.array("d", (_NAN if k is None else index.get(k, _NAN)
+                                 for k in self.keys))
+
+    def __repr__(self) -> str:
+        return f"Keys({self.levels!r})"
+
+
+def _order(v: Any) -> list[Any] | None:
+    """The categories of a pandas or polars categorical, in order (else ``None``)."""
+    cat = getattr(v, "cat", None)
+    if cat is not None and hasattr(cat, "categories"):
+        try:
+            return list(cat.categories)
+        except TypeError:
+            pass
+    dtype: Any = getattr(v, "dtype", None)
+    if dtype is not None and hasattr(dtype, "categories") and (
+            hasattr(v, "codes") or str(dtype).startswith("Enum")):
+        return list(dtype.categories)                     # pandas Categorical, polars Enum
+    if str(dtype).startswith(("Categorical", "Enum")):
+        try:                                                      # polars
+            return list(v.cat.get_categories())
+        except (AttributeError, TypeError):
+            pass
+    return None
+
+
+def is_categorical(v: Any) -> bool:
+    """Whether a position column holds categories: text, or a categorical."""
+    if isinstance(v, Keys) or _order(v) is not None:
+        return True
+    kind = getattr(getattr(v, "dtype", None), "kind", None)
+    if kind in ("U", "S"):
+        return True
+    if kind == "O" or isinstance(v, (list, tuple)):
+        items = v.tolist() if kind == "O" else v
+        some = False
+        for k in items:
+            if isinstance(k, str):
+                some = True
+            elif not (_missing(k) or (isinstance(k, float) and math.isnan(k))):
+                return False
+        return some
+    return False
+
+
+def positions(v: Any, what: str) -> Any:
+    """A position column: ``array('d')`` of numbers, or :class:`Keys` of categories."""
+    if isinstance(v, Keys):
+        return v
+    if not isinstance(v, (list, tuple, array.array)) and _shape(v) is None and not hasattr(
+            v, "to_numpy") and not hasattr(v, "dtype"):
+        v = list(v)                                   # a generator: read once
+    if is_categorical(v):
+        order = _order(v)
+        if hasattr(v, "tolist") and not isinstance(v, (list, tuple)):
+            v = v.tolist()
+        return Keys(v, order)
+    return numbers(v, what)
+
+
+def is_several(v: Any) -> bool:
+    """Whether ``v`` holds several columns (a dict, a data frame, a 2-D buffer, a list or
+    tuple of sequences)."""
+    if isinstance(v, dict) or is_frame(v):
+        return True
+    shape = _shape(v)
+    if shape is not None:
+        return len(shape) == 2
+    return isinstance(v, (list, tuple)) and bool(v) and all(
+        hasattr(x, "__len__") and not isinstance(x, (str, bytes)) for x in v)
+
+
+def named_sets(v: Any, what: str) -> tuple[list[tuple[Any, Any]], bool]:
+    """``([(name, column)], named)`` of several columns (as :func:`is_several`), columns
+    left as given (``named``: by a dict's keys or a frame's column names; else ``0, 1,
+    …``); one column otherwise."""
+    if isinstance(v, dict):
+        return list(v.items()), True
+    if is_frame(v):
+        return [(c, v[c]) for c in v.columns], True
+    shape = _shape(v)
+    if shape is not None and len(shape) == 2:
+        h, w, _, flat = grid(v)
+        return [(j, flat[j::w] if w > 1 else flat) for j in range(w)], False
+    if is_several(v):
+        return list(enumerate(v)), False
+    return [(0, v)], False
+
+
+def sets(v: Any, what: str) -> tuple[list[tuple[Any, array.array]], bool]:
+    """``([(name, numbers)], named)`` of one column or several (see the module)."""
+    items, named = named_sets(v, what)
+    one = len(items) == 1 and not named and not is_several(v)
+    return [(k, numbers(c, what if one else f"{what} {k!r}")) for k, c in items], named
+
+
+@typed("Column | Mapping[Any, Any]")
+def weights(v: Any) -> Any:
+    """A ``weights`` setting: one column, or one per set (a dict or a frame by set name,
+    else a list in set order)."""
+    if not is_several(v):
+        return numbers(v, "weights")
+    items, named = named_sets(v, "weights")
+    out = [(k, numbers(c, f"weights {k!r}")) for k, c in items]
+    return dict(out) if named else [c for _, c in out]
+
+
+#: The slot kinds of a type's ``DATA`` declaration (``{argument: kind}``, in reading
+#: order).
+KINDS = {
+    "position": "numbers or categories (x and y of the (x, y) types)",
+    "numbers": "numbers",
+    "sets": "one column or several, along x (line y, area ys)",
+    "samples": "one set of values or several; dates allowed (hist, kde, ecdf)",
+    "categories": "value sets by category, dates allowed; long data: values flat with the "
+                  "category of each in x (y when horizontal)",
+    "heights": "heights at x: x and heights, a dict {x: height}, or heights alone",
+    "slots": "numeric positions, one per group (box, violin, strip, swarm)",
+    "curve": "a constant or a column along x (fill)",
+    "wedges": "numbers, or a dict of them by name (pie)",
+    "raw": "kept as given (grids, coordinates)",
+    "cells": "kept as given, categories never read as positions (dots' long data)",
+}
+
+
+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
+
+
+def _gone(v: Any) -> bool:
+    return v is None or (isinstance(v, float) and math.isnan(v))
+
+
+def _seconds(v: Any) -> tuple[Any, Any] | None:
+    """Dates in one array, a dict or a list of arrays as seconds, with the time zone of
+    the first set (``None`` when ``v`` holds no dates)."""
+    from . import times
+    if isinstance(v, dict):
+        items = list(v.values())
+    elif isinstance(v, (list, tuple)) and v and all(
+            hasattr(x, "__len__") and not isinstance(x, str) for x in v):
+        items = list(v)
+    else:
+        return times.to_seconds(v) if times.is_time_like(v) else None
+    if not any(times.is_time_like(x) for x in items):
+        return None
+    out = [times.to_seconds(x) for x in items]
+    tz = next((z for _, z in out if z is not None), None)
+    if isinstance(v, dict):
+        return {k: s for k, (s, _) in zip(v, out)}, tz
+    return [s for s, _ in out], tz
+
+
+class Group:
+    """Rows of a call sharing a category and a set: their columns by argument."""
+
+    __slots__ = ("category", "set", "data")
+
+    def __init__(self, category: Any, set: Any, data: dict[str, Any]) -> None:
+        self.category = category
+        self.set = set
+        self.data = data
+
+    def __repr__(self) -> str:
+        return f"Group({self.category!r}, {self.set!r})"
+
+
+#: The attributes :meth:`Columns._read` sets.
+_READ = ("data", "keys", "groups", "categories", "named", "value", "counted")
+
+
+class Columns:
+    """The data arguments of one call, read once (by ``SeriesMeta``, after ``data=``
+    columns and dates) from the type's ``DATA`` declaration (:data:`KINDS`).
+
+    * ``data``: the columns by argument: ``array('d')`` (positions given as categories at
+      their index among their own), a constant, or the argument as given (``raw``);
+    * ``keys``: the position arguments given as categories (:class:`Keys`);
+    * ``groups``: the rows as :class:`Group` ``(category, set, data)``: one per set of
+      wide input (a dict, a frame, a 2-D buffer, a list of sets; never concatenated), one
+      per category of long input (grouped once by a stable argsort of the codes);
+    * ``categories``: the category of each group (``categories`` slots), of each bar
+      (``heights``), of each wedge (``wedges``, when named);
+    * ``named``: whether the sets are named (a dict's keys, a frame's columns);
+    * ``titles``: the ``data=`` column name of each argument; ``times``: the time zone of
+      each argument given as dates; ``orders``: the categorical orders of ``data=``
+      columns.
+    """
+
+    def __init__(self, kinds: dict[str, str], args: dict[str, Any], options: dict[str, Any],
+                 titles: dict[str, Any], orders: dict[str, Any], times: dict[str, Any],
+                 read: bool = True) -> None:
+        self.kinds = kinds
+        #: the arguments as read (``data=`` columns, dates as seconds), before parsing
+        self.args = args
+        self.options = options
+        self.titles = titles
+        self.orders = orders
+        self.times = times
+        if read:
+            self._read()
+
+    def __getattr__(self, name: str) -> Any:
+        if name in _READ and "kinds" in self.__dict__:     # not read yet: read on first use
+            self._read()
+            return getattr(self, name)
+        raise AttributeError(name)
+
+    def _read(self) -> None:
+        self.data: dict[str, Any] = {}
+        self.keys: dict[str, Keys] = {}
+        self.groups: list[Group] = []
+        self.categories: list[Any] = []
+        self.named = False
+        #: the argument the groups hold
+        self.value: str | None = None
+        #: ``x`` missing, the second argument's values given in its place: ``x`` is 0, 1, …
+        self.counted = False
+        names = list(self.kinds)
+        self.counted = (len(names) > 1 and names[0] == "x" and self.args.get(names[1]) is None
+                        and self.kinds[names[1]] in ("numbers", "sets"))
+        for name, kind in self.kinds.items():
+            getattr(self, "_" + kind)(name)
+
+    @classmethod
+    def build(cls, kinds: dict[str, str], given: dict[str, Any], kw: dict[str, Any],
+              titles: dict[str, Any], orders: dict[str, Any], times: dict[str, Any],
+              items: str = "sets", read: bool = True) -> "Columns":
+        """The columns of a call: ``given`` the data arguments (defaults included), ``kw``
+        its other keywords (long data takes its categories from ``x``, ``y`` when
+        horizontal); ``items`` names a type's groups in errors.  ``read=False``: the
+        arguments, read on first use (a call split by ``group=`` or by levels)."""
+        args = dict(given)
+        options = {"categories": kw.get("categories"), "items": items}
+        for name, kind in kinds.items():
+            if kind in ("samples", "categories") and args.get(name) is not None:
+                converted = _seconds(args[name])
+                if converted is not None:
+                    args[name], times[name] = converted
+            if kind == "categories":
+                key = "y" if kw.get("orientation") == "horizontal" else "x"
+                cats, vals = kw.get(key), args.get(name)
+                n = _length(vals)
+                if (cats is not None and n is not None and not isinstance(vals, dict)
+                        and _length(cats) == n):
+                    args[key] = kw.pop(key)
+                    options["long"] = key
+                    if key in titles:
+                        titles["categories"] = titles.pop(key)
+        return cls(kinds, args, options, titles, orders, times, read)
+
+    def take(self, rows: list[int], n: int) -> "Columns":
+        """The columns of ``rows`` (of ``n``: arguments of another length are kept whole),
+        read again: categories in their order of appearance among the rows."""
+        from .tables import take
+        args = {k: (take(v, rows) if _length(v) == n else v) for k, v in self.args.items()}
+        return Columns(self.kinds, args, self.options, dict(self.titles), {}, dict(self.times))
+
+    def __getitem__(self, name: str) -> Any:
+        return self.data[name]
+
+    def __contains__(self, name: str) -> bool:
+        return name in self.data
+
+    def get(self, name: str, default: Any = None) -> Any:
+        return self.data.get(name, default)
+
+    @property
+    def sets(self) -> list[tuple[Any, Any]]:
+        """``(set, column)`` of each group."""
+        key = self.value or ""
+        return [(g.set, g.data[key]) for g in self.groups]
+
+    @property
+    def values(self) -> list[Any]:
+        """The column of each group."""
+        key = self.value or ""
+        return [g.data[key] for g in self.groups]
+
+    def __repr__(self) -> str:
+        return f"Columns({list(self.data)}, groups={len(self.groups)})"
+
+    # -- slot kinds (see KINDS) ---------------------------------------------------------
+
+    def _pair(self, name: str, c: Any) -> None:
+        """A column of an (x, y) call: the other's length checked."""
+        if isinstance(c, Keys):
+            self.keys[name] = c
+            c = c.own
+        self.data[name] = c
+        other = {"x": "y", "y": "x"}.get(name)
+        if other in self.data and self.kinds.get(other) in ("position", "numbers"):
+            x, y = (c, self.data[other]) if name == "x" else (self.data[other], c)
+            if len(x) != len(y):
+                raise ValueError(f"x has {len(x)} values and y {len(y)}")
+
+    def _position(self, name: str) -> None:
+        if name == "x" and (self.counted or "heights" in self.kinds.values()):
+            return                                      # read with the next argument
+        self._pair(name, positions(self.args[name], name))
+
+    def _numbers(self, name: str) -> None:
+        if self.counted and name != "x":
+            y = numbers(self.args["x"], name)
+            self._pair("x", numbers(range(len(y)), "x"))
+        else:
+            y = numbers(self.args[name], name)
+        self._pair(name, y)
+
+    def _sets(self, name: str) -> None:
+        items, self.named = sets(self.args["x" if self.counted else name], name)
+        if self.counted:
+            self.data["x"] = array.array("d", range(len(items[0][1])))
+        x = self.data["x"]
+        what = name if name == "y" else "a layer"
+        for _, r in items:
+            if len(r) != len(x):
+                raise ValueError(f"x has {len(x)} values and {what} {len(r)}")
+        self.value = name
+        self.groups = [Group(None, k, {name: v}) for k, v in items]
+
+    def _samples(self, name: str) -> None:
+        items, self.named = sets(self.args[name], name)
+        self.value = name
+        self.groups = [Group(None, k, {name: v}) for k, v in items]
+
+    def _categories(self, name: str) -> None:
+        values = self.args[name]
+        key = self.options.get("long")
+        if key is not None:
+            values = self._long(numbers(values, name), self.args[key], self.orders.get(key))
+        items, self.named = sets(values, name)
+        given = self.options.get("categories")
+        what = self.options["items"]
+        if not self.named and given is not None:
+            names = list(given)
+            if is_several(values):
+                if len(names) != len(items):
+                    raise ValueError(f"{len(names)} categories for {len(items)} {what}")
+                items = list(zip(names, (v for _, v in items)))
+            elif names:
+                items = [(names[0], items[0][1])]
+        self.value = name
+        self.groups = [Group(k, None, {name: v}) for k, v in items]
+        self.categories = [k for k, _ in items]
+
+    @staticmethod
+    def _long(values: array.array, cats: Any, order: Any) -> dict[Any, array.array]:
+        """Flat values by category (``order`` first, then first appearance; rows of a
+        missing category dropped, categories without rows left out), each in row order."""
+        from bisect import bisect_left, bisect_right
+        from . import kernels
+        index: dict[Any, int] = {}
+        for c in order or []:
+            index.setdefault(c, len(index))
+        lookup: dict[Any, float] = {}
+        for c in dict.fromkeys(cats):                    # first appearance
+            if _gone(c):
+                lookup[c] = -1.0
+            elif c not in index:
+                index[c] = len(index)
+        lookup.update((c, float(k)) for c, k in index.items())
+        codes = array.array("d", map(lookup.__getitem__, cats))
+        gone = codes.count(-1.0)
+        if len(index) == 1 and not gone:
+            return {next(iter(index)): values} if len(values) else {}
+        rows = kernels.argsort(codes)[gone:]             # stable; missing (-1) first
+        ordered = kernels.take(values, rows) if len(rows) else array.array("d")
+        sorted_codes = kernels.take(codes, rows) if len(rows) else array.array("d")
+        out: dict[Any, array.array] = {}
+        for c, k in index.items():
+            lo = bisect_left(sorted_codes, k)
+            hi = bisect_right(sorted_codes, k, lo)
+            if hi > lo:
+                out[c] = ordered[lo:hi]
+        return out
+
+    def _slots(self, name: str) -> None:
+        p = self.args.get(name)
+        if p is not None:
+            p = list(p)
+            for v in p:
+                if isinstance(v, (str, bytes)):
+                    raise ValueError(f"{name}: expected numbers (one position per group), "
+                                     f"got {v!r}; categories: a dict {{category: values}} "
+                                     f"or categories=")
+        self.data[name] = None if p is None else [float(v) for v in p]
+        if p is not None and len(self.data[name]) != len(self.groups):
+            raise ValueError(f"{len(self.data[name])} positions for {len(self.groups)} "
+                             f"{self.options['items']}")
+
+    def _heights(self, name: str) -> None:
+        x, h = self.args["x"], self.args[name]
+        if h is None:
+            if isinstance(x, dict):
+                keys = list(x.keys())
+                vals = numbers(list(x.values()), "heights")
+                categorical = all(isinstance(k, str) for k in keys)
+            else:
+                vals = numbers(x, "heights")
+                keys = [float(i) for i in range(len(vals))]
+                categorical = False
+        else:
+            vals = numbers(h, "heights")
+            pos = positions(x, "x")
+            if len(pos) != len(vals):
+                raise ValueError(f"{len(pos)} positions for {len(vals)} values")
+            categorical = isinstance(pos, Keys)
+            if isinstance(pos, Keys):
+                self.keys["x"] = pos
+                keep = [i for i, k in enumerate(pos.keys) if k is not None]  # missing bars
+                keys = [pos.keys[i] for i in keep]
+                if len(keep) < len(pos):
+                    vals = array.array("d", [vals[i] for i in keep])
+            else:
+                keys = list(pos)
+        self.categories = keys
+        self.data[name] = vals
+        self.data["x"] = None if categorical else [float(k) for k in keys]
+
+    def _curve(self, name: str) -> None:
+        if name == "y1":
+            y1, y2 = self.args["y1"], self.args.get("y2")
+            if is_several(y1) and isinstance(y2, (int, float)) and y2 == 0:
+                items, _ = named_sets(y1, "y1")
+                if len(items) == 2:
+                    self.args = {**self.args, "y1": items[0][1], "y2": items[1][1]}
+        v = self.args[name]
+        if isinstance(v, (int, float)) and not isinstance(v, bool):
+            self.data[name] = float(v)
+            return
+        c = numbers(v, name)
+        n = len(self.data["x"])
+        if len(c) != n:
+            raise ValueError(f"{name} has {len(c)} values for {n} points")
+        self.data[name] = c
+
+    def _wedges(self, name: str) -> None:
+        v = self.args[name]
+        if isinstance(v, dict):
+            self.categories = [str(k) for k in v]
+            v = list(v.values())
+        self.data[name] = doubles(v)
+
+    def _raw(self, name: str) -> None:
+        self.data[name] = self.args.get(name)
+
+    _cells = _raw