plot-io 0.1.0__py3-none-any.whl

plot_io/maps/__init__.py

@@ -0,0 +1,365 @@
+"""Color and size maps.
+
+A :class:`ColorMap` is a color gradient plus the mapping of values onto it (limits and
+scale); a :class:`SizeMap` maps values to marker diameters.  Values stay on the series
+(``color_values``, ``size_values``), so one map can serve several series.
+"""
+
+import math
+from collections.abc import Sequence
+from typing import Any
+
+from pdf_io.colors import to_rgba
+
+from .. import _backend
+from ..limits import check_limits
+from ..scales import NUMERIC_SCALES
+from ..settings import one_of
+from .data import DATA
+
+__all__ = ["ColorMap", "SizeMap", "register_color_map", "color_map_names"]
+
+RGBA = tuple[float, float, float, float]
+_check_scale = one_of(*NUMERIC_SCALES)
+
+
+# ---------------------------------------------------------------------------
+# Oklab
+# ---------------------------------------------------------------------------
+
+def _to_linear(c: float) -> float:
+    return c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4
+
+
+def _from_linear(c: float) -> float:
+    c = min(1.0, max(0.0, c))
+    return 12.92 * c if c <= 0.0031308 else 1.055 * c ** (1 / 2.4) - 0.055
+
+
+def srgb_to_oklab(r: float, g: float, b: float) -> tuple[float, float, float]:
+    r, g, b = _to_linear(r), _to_linear(g), _to_linear(b)
+    lc = (0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b) ** (1 / 3)
+    mc = (0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b) ** (1 / 3)
+    sc = (0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b) ** (1 / 3)
+    return (0.2104542553 * lc + 0.7936177850 * mc - 0.0040720468 * sc,
+            1.9779984951 * lc - 2.4285922050 * mc + 0.4505937099 * sc,
+            0.0259040371 * lc + 0.7827717662 * mc - 0.8086757660 * sc)
+
+
+def oklab_to_srgb(L: float, a: float, b: float) -> tuple[float, float, float]:
+    lc = (L + 0.3963377774 * a + 0.2158037573 * b) ** 3
+    mc = (L - 0.1055613458 * a - 0.0638541728 * b) ** 3
+    sc = (L - 0.0894841775 * a - 1.2914855480 * b) ** 3
+    return (_from_linear(4.0767416621 * lc - 3.3077115913 * mc + 0.2309699292 * sc),
+            _from_linear(-1.2684380046 * lc + 2.6097574011 * mc - 0.3413193965 * sc),
+            _from_linear(-0.0041960863 * lc - 0.7034186147 * mc + 1.7076147010 * sc))
+
+
+# ---------------------------------------------------------------------------
+# Color maps
+# ---------------------------------------------------------------------------
+
+def _rgba(color: Any) -> RGBA:
+    try:
+        c = to_rgba(color)
+    except (TypeError, ValueError):
+        c = None
+    if c is None:
+        raise ValueError(f"invalid color {color!r}")
+    return (float(c[0]), float(c[1]), float(c[2]), float(c[3]) if len(c) > 3 else 1.0)
+
+
+def _hex_stops(data: str) -> list[tuple[float, RGBA]]:
+    colors = [(int(data[i:i + 2], 16) / 255, int(data[i + 2:i + 4], 16) / 255,
+               int(data[i + 4:i + 6], 16) / 255, 1.0) for i in range(0, len(data), 6)]
+    n = len(colors) - 1
+    return [(i / n, c) for i, c in enumerate(colors)]
+
+
+#: Registered maps: name → (stops, space).
+_REGISTRY: dict[str, tuple[list[tuple[float, RGBA]], str]] = {
+    "grey": ([(0.0, (0.0, 0.0, 0.0, 1.0)), (1.0, (1.0, 1.0, 1.0, 1.0))], "srgb"),
+}
+_REGISTRY.update((name, (_hex_stops("".join(d)), "srgb")) for name, d in DATA.items())
+#: ColorBrewer's YlGnBu and Blues (Cynthia Brewer, Apache-2.0): Vega-Lite's heatmap and ramp
+_REGISTRY["ylgnbu"] = (_hex_stops("ffffd9edf8b1c7e9b47fcdbb41b6c41d91c0225ea8253494081d58"),
+                       "srgb")
+_REGISTRY["blues"] = (_hex_stops("f7fbffdeebf7c6dbef9ecae16baed64292c62171b508519c08306b"),
+                      "srgb")
+#: ggplot2's continuous default (``scale_colour_gradient``), interpolated perceptually
+_REGISTRY["ggplot2"] = ([(0.0, (0x13 / 255, 0x2B / 255, 0x43 / 255, 1.0)),
+                         (1.0, (0x56 / 255, 0xB1 / 255, 0xF7 / 255, 1.0))], "oklab")
+_ALIASES = {"gray": "grey"}
+
+
+def color_map_names() -> list[str]:
+    return sorted(_REGISTRY)
+
+
+def register_color_map(name: str, color_map: Any) -> None:
+    """Register a map under ``name`` (also usable reversed, as ``name + "_r"``)."""
+    cm = color_map if isinstance(color_map, ColorMap) else ColorMap(color_map)
+    _REGISTRY[name.strip().lower()] = (list(cm.stops), cm.space)
+
+
+def _lookup(name: str) -> tuple[list[tuple[float, RGBA]], str, str]:
+    key = name.strip().lower()
+    reverse = key.endswith("_r")
+    base = key[:-2] if reverse else key
+    base = _ALIASES.get(base, base)
+    if base not in _REGISTRY:
+        raise ValueError(f"unknown color map {name!r} (bundled: {', '.join(color_map_names())})")
+    stops, space = _REGISTRY[base]
+    if reverse:
+        stops = [(1.0 - p, c) for p, c in reversed(stops)]
+    return stops, space, key
+
+
+class ColorMap:
+    """A color gradient and the mapping of values onto it.
+
+    ``colors``: a registered name (``"viridis"``, ``"viridis_r"``), a list of colors
+    (evenly spaced, or at ``positions``), another ``ColorMap`` (copied) or a function
+    ``t -> color``.  ``limits`` and ``scale`` use the axis syntax; ``"auto"`` limits cover
+    every series using the map.  ``under``/``over``: colors below and above the limits
+    (default: the end colors); ``bad``: color of missing values (default: transparent).
+    ``space``: interpolation between the given colors, ``"oklab"`` or ``"srgb"``.
+    """
+
+    space: str
+    name: str | None
+    under: RGBA | None
+    over: RGBA | None
+    bad: RGBA | None
+
+    def __init__(self, colors: Any = "viridis", positions: Sequence[float] | None = None,
+                 limits: Any = "auto", scale: str = "linear", under: Any = None,
+                 over: Any = None, bad: Any = None, space: str | None = None,
+                 name: str | None = None) -> None:
+        self.interpolation = "linear"
+        default_space = "oklab"
+        stops: list[tuple[float, RGBA]]
+        if isinstance(colors, ColorMap):
+            stops, default_space = list(colors.stops), colors.space
+            self.interpolation = colors.interpolation
+            name = name or colors.name
+            under = colors.under if under is None else under
+            over = colors.over if over is None else over
+            bad = colors.bad if bad is None else bad
+        elif isinstance(colors, str):
+            stops, default_space, key = _lookup(colors)
+            name = name or key
+        elif callable(colors):
+            stops = [(i / 255, _rgba(colors(i / 255))) for i in range(256)]
+            default_space = "srgb"
+        else:
+            cols = [_rgba(c) for c in colors]
+            if not cols:
+                raise ValueError("a color map needs at least one color")
+            if positions is None:
+                n = max(len(cols) - 1, 1)
+                stops = [(i / n, c) for i, c in enumerate(cols)]
+            else:
+                pos = [float(p) for p in positions]
+                if len(pos) != len(cols):
+                    raise ValueError("give one position per color")
+                if any(b < a for a, b in zip(pos, pos[1:])) or pos[0] < 0 or pos[-1] > 1:
+                    raise ValueError("positions must increase within [0, 1]")
+                stops = list(zip(pos, cols))
+        if len(stops) == 1:
+            stops = [(0.0, stops[0][1]), (1.0, stops[0][1])]
+        self.stops: list[tuple[float, RGBA]] = stops
+        self.space = one_of("oklab", "srgb")(space) if space is not None else default_space
+        self.limits = check_limits(limits)
+        self.scale = _check_scale(scale)
+        self.under = None if under is None else _rgba(under)
+        self.over = None if over is None else _rgba(over)
+        self.bad = None if bad is None else _rgba(bad)
+        self.name = name
+        self._lab = ([srgb_to_oklab(*c[:3]) for _, c in stops] if self.space == "oklab"
+                     else None)
+        self._pdf: Any = None
+
+    def __repr__(self) -> str:
+        return f"ColorMap({self.name or f'{len(self.stops)} colors'})"
+
+    def _copy(self, **changes: Any) -> "ColorMap":
+        args = dict(limits=self.limits, scale=self.scale, under=self.under, over=self.over,
+                    bad=self.bad, space=self.space, name=self.name)
+        args.update(changes)
+        stops = changes.pop("stops", None)
+        args.pop("stops", None)
+        cm = ColorMap(self, **args)
+        cm.name = args["name"]
+        if stops is not None:
+            cm.stops = stops
+            cm._lab = ([srgb_to_oklab(*c[:3]) for _, c in stops] if cm.space == "oklab"
+                       else None)
+        return cm
+
+    def __call__(self, t: float) -> RGBA:
+        """Color at ``t`` in [0, 1] (``under``/``over`` outside, ``bad`` for NaN)."""
+        t = float(t)
+        if t != t:
+            return self.bad or (0.0, 0.0, 0.0, 0.0)
+        if t < 0:
+            return self.under or self.stops[0][1]
+        if t > 1:
+            return self.over or self.stops[-1][1]
+        stops = self.stops
+        if self.interpolation == "nearest":
+            for i in range(len(stops) - 1, -1, -1):
+                if stops[i][0] <= t:
+                    return stops[i][1]
+            return stops[0][1]
+        if t <= stops[0][0]:
+            return stops[0][1]
+        if t >= stops[-1][0]:
+            return stops[-1][1]
+        lo, hi = 0, len(stops) - 1
+        while hi - lo > 1:
+            mid = (lo + hi) // 2
+            if stops[mid][0] <= t:
+                lo = mid
+            else:
+                hi = mid
+        (p0, c0), (p1, c1) = stops[lo], stops[hi]
+        f = (t - p0) / (p1 - p0) if p1 > p0 else 1.0
+        alpha = c0[3] + (c1[3] - c0[3]) * f
+        if self._lab is not None:
+            a, b = self._lab[lo], self._lab[hi]
+            r, g, bl = oklab_to_srgb(*(x + (y - x) * f for x, y in zip(a, b)))
+            return (r, g, bl, alpha)
+        return (c0[0] + (c1[0] - c0[0]) * f, c0[1] + (c1[1] - c0[1]) * f,
+                c0[2] + (c1[2] - c0[2]) * f, alpha)
+
+    def table(self, n: int = 256) -> list[RGBA]:
+        return [self(i / (n - 1)) for i in range(n)] if n > 1 else [self(0.0)]
+
+    def colors_at(self, ts: Any) -> Any:
+        """Colors at many positions in [0, 1] (an ``(n, 4)`` buffer, or a list)."""
+        return _backend.colors_at(self.to_pdf_io(), ts)
+
+    def reversed(self) -> "ColorMap":
+        stops = [(1.0 - p, c) for p, c in reversed(self.stops)]
+        name = None
+        if self.name:
+            name = self.name[:-2] if self.name.endswith("_r") else self.name + "_r"
+        return self._copy(stops=stops, under=self.over, over=self.under, name=name)
+
+    def resampled(self, n: int) -> "ColorMap":
+        """``n`` flat steps (a discrete map)."""
+        if n < 1:
+            raise ValueError("resampled() needs at least one step")
+        cols = [self((i + 0.5) / n) for i in range(n)]
+        cm = self._copy(stops=[(i / n, c) for i, c in enumerate(cols)], space="srgb")
+        cm.interpolation = "nearest"
+        return cm
+
+    def truncated(self, lo: float, hi: float) -> "ColorMap":
+        """The part between ``lo`` and ``hi``, stretched to [0, 1]."""
+        if not 0 <= lo < hi <= 1:
+            raise ValueError("truncated() needs 0 <= lo < hi <= 1")
+        n = 256
+        cols = [self(lo + (hi - lo) * i / (n - 1)) for i in range(n)]
+        return self._copy(stops=[(i / (n - 1), c) for i, c in enumerate(cols)], space="srgb",
+                          name=None)
+
+    def with_extremes(self, under: Any = None, over: Any = None, bad: Any = None) -> "ColorMap":
+        return self._copy(under=under if under is not None else self.under,
+                          over=over if over is not None else self.over,
+                          bad=bad if bad is not None else self.bad)
+
+    def pdf_stops(self) -> list[tuple[float, RGBA]]:
+        """Stops for pdf_io (sRGB interpolation): Oklab maps as 256 samples."""
+        if self.space == "oklab" and len(self.stops) < 256 and self.interpolation == "linear":
+            return [(i / 255, self(i / 255)) for i in range(256)]
+        return list(self.stops)
+
+    def to_pdf_io(self) -> Any:
+        """The equivalent pdf_io color map."""
+        if self._pdf is None:
+            self._pdf = _backend.pdf_color_map(self.pdf_stops(), self.interpolation,
+                                               self.under, self.over, self.bad)
+        return self._pdf
+
+
+def check_cmap(v: Any) -> Any:
+    """A ``ColorMap``, or anything its ``colors`` accepts (validated)."""
+    if isinstance(v, ColorMap):
+        return v
+    ColorMap(v)
+    return v
+
+
+def as_color_map(value: Any) -> ColorMap:
+    return value if isinstance(value, ColorMap) else ColorMap(value)
+
+
+# ---------------------------------------------------------------------------
+# Size maps
+# ---------------------------------------------------------------------------
+
+def _size(v: Any) -> float:
+    if isinstance(v, bool) or not isinstance(v, (int, float)) or not v >= 0:
+        raise ValueError(f"expected a size >= 0, got {v!r}")
+    return float(v)
+
+
+class SizeMap:
+    """Values to marker diameters (pt).
+
+    ``sizes``: the diameters at the two limits.  ``by="area"`` moves the marker area
+    linearly between the end areas; ``by="diameter"`` moves the diameter.  ``under`` and
+    ``over``: sizes below and above the limits (default: the end sizes; ``"none"`` leaves
+    those points out); ``bad``: size of points with a missing value (default: left out).
+    ``label``: title of its size key in legends (default: the ``size`` column name).
+    """
+
+    def __init__(self, sizes: Sequence[float] = (2, 10), limits: Any = "auto",
+                 scale: str = "linear", by: str = "area", under: Any = None,
+                 over: Any = None, bad: Any = None, label: Any = None) -> None:
+        self.label = label
+        if len(sizes) != 2:
+            raise ValueError("sizes must be a pair (smallest, largest)")
+        self.sizes = (_size(sizes[0]), _size(sizes[1]))
+        self.limits = check_limits(limits)
+        self.scale = _check_scale(scale)
+        self.by = one_of("area", "diameter")(by)
+        self.under = self._extreme(under)
+        self.over = self._extreme(over)
+        self.bad = self._extreme(bad)
+
+    @staticmethod
+    def _extreme(v: Any) -> Any:
+        if v is None:
+            return None
+        if isinstance(v, str) and v.lower() == "none":
+            return math.nan
+        return _size(v)
+
+    def __repr__(self) -> str:
+        return f"SizeMap({self.sizes[0]:g}–{self.sizes[1]:g} pt, by {self.by})"
+
+    def map(self, values: Any, flo: float, fhi: float, function: str | None) -> Any:
+        """Diameters for ``values``, with ``flo``/``fhi`` the limits after the scale
+        function."""
+        d0, d1 = self.sizes
+        span = fhi - flo
+        nan = math.nan
+        under = d0 if self.under is None else self.under
+        over = d1 if self.over is None else self.over
+        bad = nan if self.bad is None else self.bad
+        if self.by == "diameter":
+            k = (d1 - d0) / span if span else 0.0
+            lo, hi = (d0, d1) if d1 >= d0 else (d1, d0)
+            u, o = (under, over) if d1 >= d0 else (over, under)
+            return _backend.scaled(values, k, d0 - k * flo, function, low=lo, high=hi,
+                                   under=u, over=o, nan=bad)
+        a0, a1 = d0 * d0, d1 * d1
+        k = (a1 - a0) / span if span else 0.0
+        lo, hi = (a0, a1) if a1 >= a0 else (a1, a0)
+        u, o = (under * under, over * over) if a1 >= a0 else (over * over, under * under)
+        areas = _backend.scaled(values, k, a0 - k * flo, function, low=lo, high=hi,
+                                under=u, over=o, nan=bad * bad)
+        return _backend.scaled(areas, 1.0, 0.0, "sqrt")