plot-io 0.12.0__py3-none-win_arm64.whl

plot_io/presets/contour.py

@@ -0,0 +1,353 @@
+"""Contours (stat ``contour``): filled bands (``Polygons``), iso-lines (``Path``) and level
+labels on the lines (``Text``)."""
+
+import array
+import math
+from typing import Any
+
+from vec_io.arrays import doubles
+
+from .._methods import ContourMixin
+from .. import groups
+from .. import kernels
+from ..color.values import with_alpha
+from ..layers import stats
+from ..layers.base import Series, common
+from ..layers.core import Call, Layer, Table
+from ..layers.grammar import entry
+from ..layers.paint import Paint, role
+from ..labels import drop, quad
+from ..layers.geoms import GEOMS
+from ..lines import dash_of
+from ..mapped import color_axis_of, unit_of
+from ..numbers import NumberStyle, check_format, format_numbers
+from ..settings import Setting, boolean, either, literal, number, one_of, table
+from ..text import font_of, measure
+from ..ticks import nice_step
+from ..transform import settings_both as coord_settings
+from ._util import columns_meta, given_colors, map_or_color, placed_text
+from .heatmap import Coords
+
+__all__ = ["Contour"]
+
+
+def _levels(v: Any) -> Any:
+    if isinstance(v, int) and not isinstance(v, bool):
+        return number(lo=1, integer=True)(v)
+    vals = sorted(float(x) for x in v)
+    if len(vals) < 1:
+        raise ValueError("give at least one level")
+    return vals
+
+
+_COLOR_MAP_DOC = ("The colors of the levels: a map (a qualitative one's colors taken in "
+                  "turn by the lines), one color for every line, `\"auto\"` or a kind "
+                  "(as a series' `color_map`).")
+
+
+class Contour(ContourMixin, Series):
+    """Contour lines and/or filled bands of ``values`` (rows first) at nodes ``x``/``y``
+    (centres; edges are reduced to their midpoints).
+
+    ``lines`` draws iso-lines and ``fill`` filled bands.  ``levels``: a number of nice values
+    over the data, or the values.  Lines take the map's colors: of a qualitative
+    ``color_map`` in turn, else at their level; ``color_map`` as one color draws them in
+    that color; ``negative_style`` dashes
+    negative levels drawn in one color; ``value_labels`` writes the levels on the lines.
+    Bands are colored by their middle value.
+
+    Example::
+
+        import math
+        grid = [[math.sin(i / 5) * math.cos(j / 6) for j in range(40)] for i in range(30)]
+        p = plt.Plot(width=240, height=160)
+        p.contour(grid, value_labels=True)
+    """
+
+    SETTINGS = table(
+        **coord_settings(),
+        levels=Setting(_levels, doc="A number of levels (nice values over the data) or "
+                                        "the level values."),
+        **groups.FillStyle.embed(checks={"fill": lambda v: groups.fill_value(v, colors=False)},
+                                 docs={"fill": "Fill the bands between levels: `True` or "
+                                               "`False`.",
+                                       "fill_opacity": "Opacity of the bands, times "
+                                                       "`opacity`."}),
+        lines=Setting(boolean, doc="Draw the iso-lines."),
+        **{**common("color_map"),
+           "color_map": Setting(map_or_color, doc=_COLOR_MAP_DOC)},
+        **groups.Stroke.embed(only=("line_width", "line_style"),
+                              inherited=("line_width",),
+                              docs={"line_width": "Line width; default 0.6 × the plot's.",
+                                    "line_style": "Line style."}),
+        negative_style=Setting(either(literal(None), one_of("solid", "dashed", "dotted",
+                                                            "dashdot")),
+                               doc="Line style of negative levels drawn with a single color "
+                                   "(matplotlib: dashed)."),
+        opacity=Setting(number(0, 1), doc="Opacity."),
+        value_labels=Setting(boolean, doc="Write the level values on the lines (a label over an "
+                                                  "earlier one is left out)."),
+        value_label_format=Setting(check_format, doc="Number format of the value labels."),
+        **groups.TextStyle.embed("value_label_",
+                                 only=("font_size",),
+                                 docs={"font_size": "Label font size."}),
+        **common("name", "z", "rasterize", "simplify", "font_size", "font_family", "text_color"),
+        **groups.NumberFormat.embed(),
+    )
+
+    DATA = {"values": "raw", "x": "raw", "y": "raw"}
+
+    def __init__(self, values: Any, x: Any = None, y: Any = None, plot: Any = None,
+                 **settings: Any) -> None:
+        values, x, y = self.cols["values"], self.cols["x"], self.cols["y"]
+        g = kernels.grid(values)
+        if g is not None:
+            self.h, self.w, k, self.flat = g
+            if k != 1:
+                raise ValueError("contour values must be a 2-D grid of numbers")
+        else:
+            rows = [list(r) for r in values]
+            self.h, self.w = len(rows), len(rows[0]) if rows else 0
+            if any(len(r) != self.w for r in rows):
+                raise ValueError("contour rows must have the same length")
+            self.flat = doubles([v for r in rows for v in r])
+        if self.h < 2 or self.w < 2:
+            raise ValueError("contours need a grid of at least 2 × 2 values")
+        self.xn = self._nodes(x, self.w, "x")
+        self.yn = self._nodes(y, self.h, "y")
+        super().__init__(plot=plot, **settings)
+
+    def _bar_source(self) -> Any:
+        return self
+
+    @staticmethod
+    def _nodes(given: Any, n: int, what: str) -> array.array:
+        if given is None:
+            return array.array("d", (float(i) for i in range(n)))
+        c = Coords(given, n, what)
+        if c.kind == "category":
+            raise ValueError("contour coordinates must be numbers")
+        v = c.values
+        if c.kind == "edges":
+            return array.array("d", ((a + b) / 2 for a, b in zip(v, v[1:])))
+        return array.array("d", v)
+
+    def _scales(self) -> tuple[Any, Any]:
+        return stats.axis_scale(self.plot.x_axis), stats.axis_scale(self.plot.y_axis)
+
+    # -- levels and map ---------------------------------------------------------------
+
+    def level_values(self) -> list[float]:
+        lv = self.resolve("levels")
+        if isinstance(lv, list):
+            return lv
+        if isinstance(lv, bool) or not isinstance(lv, (int, float)):
+            raise ValueError(f"{self.kind} setting 'levels': expected a number of levels or "
+                             f"the level values, got {lv!r}")
+        n, lo, hi = kernels.extent(self.flat, None)
+        if not n:
+            return [0.0]
+        if hi == lo:
+            return [lo]
+        step = nice_step((hi - lo) / lv)
+        start = math.floor(lo / step + 1e-9) * step
+        end = math.ceil(hi / step - 1e-9) * step
+        k = int(round((end - start) / step))
+        return [round(start + i * step, 12) for i in range(k + 1)]
+
+    def _filled(self) -> bool:
+        return self.resolve("fill") is not False
+
+    def _fill_opacity(self) -> float:
+        return groups.fill_parts(self)[1]
+
+    def _grid_data(self) -> tuple[Any, int, int, Any, Any]:
+        """``(flat, w, h, fx, fy)``: the grid and its nodes after the axis scales."""
+        sx, sy = self._scales()
+        fx = array.array("d", (sx.f(v) for v in self.xn))
+        fy = array.array("d", (sy.f(v) for v in self.yn))
+        return self.flat, self.w, self.h, fx, fy
+
+    def _bands(self, lv: list[float]) -> list[tuple[float, float, float]]:
+        out = []
+        for lo, hi in zip(lv, lv[1:]):
+            top = hi if hi != lv[-1] else math.nextafter(hi, math.inf)
+            out.append((lo, hi, top))
+        return out
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        """``{"levels", "lines": {level: [(x, y, closed)]}, "bands": [(lo, hi, (xs, ys,
+        ring_starts))]}``, coordinates after the axis scale functions."""
+        sx, sy = self._scales()
+        return self.cached(("stats", stats.scale_key(sx), stats.scale_key(sy)), self._compute,
+                           stats.READS["contour"], kind="stats")
+
+    def _compute(self) -> dict[str, Any]:
+        flat, w, h, fx, fy = self._grid_data()
+        lv = self.level_values()
+        bands = self._bands(lv) if self._filled() else []
+        return stats.contour(flat, w, h, fx, fy, lv, bool(self.resolve("lines")), bands)
+
+    def _given(self) -> Any:
+        """The lines' colors in place of the map's at their levels: one color, a qualitative
+        map's colors (taken in turn), or ``None``."""
+        return given_colors(self, len(self.level_values()), ("qualitative",))
+
+    def _level_color(self, ctx: Any, i: int, level: float) -> Any:
+        given = self._given()
+        if given is not None:
+            return ctx.color(given[i % len(given)] if isinstance(given, list) else given)
+        axis = color_axis_of(self)
+        return axis.color_map()(unit_of(axis)(level))
+
+    # -- the preset --------------------------------------------------------------------
+
+    def layers(self) -> Call:
+        call = Call(self)
+        st = self.stats
+        opacity = self.resolve("opacity")
+        key = id(self)
+        if st["bands"]:
+            band_opacity = opacity * self._fill_opacity()
+
+            call.layers.append(Layer(GEOMS["polygons"],
+                                     Table(ring_sets=[r for _, _, r in st["bands"]]),
+                                     paint=Paint(self, "contour.bands",
+                                                 opacity=band_opacity),
+                                     stat="contour", key=False))
+        lw = self.resolve("line_width")
+        style = self.resolve("line_style")
+        neg = self.resolve("negative_style")
+        given = self._given()
+        single = given is not None and not isinstance(given, list)
+        drawn = []                                  # (level index, level, polylines)
+        for i, level in enumerate(st["levels"]):
+            polylines = st["lines"].get(level) or []
+            if not polylines or not lw:
+                continue
+            drawn.append((i, level, polylines))
+            ls = neg if (neg is not None and single and level < 0) else style
+
+            call.layers.append(Layer(GEOMS["path"], Table(x=[p[0] for p in polylines],
+                                                          y=[p[1] for p in polylines]),
+                                     paint=Paint(self, "contour.line", i=i, level=level,
+                                                 dash=dash_of(ls, lw), opacity=opacity),
+                                     stat="contour", key=False, lists=True,
+                                     f_space=True))
+        if self.resolve("value_labels") and drawn:
+            size = self.resolve("value_label_font_size")
+            call.layers.append(placed_text(self, key, lambda ctx: self._label_items(ctx, drawn),
+                                           lambda ctx: font_of(self, size),
+                                           background=lambda ctx: self._background(ctx)))
+        self._occupancy(call)
+        filled = self._filled()
+        glyph = Layer(GEOMS["polygons" if filled else "path"], Table(),
+                      paint=Paint(self, "contour.key"), glyph="fill")
+        name = self.resolve("name")
+        if name is not None:
+            call.entries.append(entry(self, name, [glyph], stroke="_key_color()",
+                                      fill="_key_color()" if filled else None))
+        lv = self.level_values()
+        call.color_count = self._color_count()
+        if self._mapped_source():
+            call.color_sources = [(color_axis_of(self), array.array("d", lv))]
+        self._extent(call)
+        call.exact = {"x", "y"}
+        columns_meta(call, self._roles())
+        return call
+
+    def _color_count(self) -> int:
+        return 0
+
+    def _key_color(self, ctx: Any) -> Any:
+        """The legend's color: the middle level's."""
+        lv = self.level_values()
+        return self._level_color(ctx, len(lv) // 2, lv[len(lv) // 2])
+
+    def _mapped_source(self) -> bool:
+        return not (self._given() is not None and not self._filled())
+
+    def _roles(self) -> dict[str, str]:
+        return {"values": "x"}
+
+    def _occupancy(self, call: Call) -> None:
+        for la in call.layers:
+            la.occupy = "all" if not la.overlay else False
+
+    def _extent(self, call: Call) -> None:
+        call.extra = {"x": self.xn, "y": self.yn}
+
+    def _background(self, ctx: Any) -> Any:
+        fig = self.plot.figure
+        return next((c for c in (ctx.color(self.plot.resolve("background")),
+                                 ctx.color(fig.resolve("background")), "white")
+                     if c is not None))
+
+    def _label_items(self, ctx: Any, drawn: list[Any]) -> list[Any]:
+        canvas = ctx.canvas
+        xm, ym = ctx.xmap, ctx.ymap
+        levels = [lv for _, lv, _ in drawn]
+        texts = format_numbers(levels, self.resolve("value_label_format"), NumberStyle.of(self))
+        size = self.resolve("value_label_font_size")
+        opacity = self.resolve("opacity")
+        x0, y0, x1, y1 = ctx.rect
+        out, quads = [], []
+        for (i, level, polylines), text in zip(drawn, texts):
+            c = with_alpha(self._level_color(ctx, i, level), opacity)
+            font = font_of(self, size, color=c)
+            w, a, d = measure(canvas, text, font)
+            for p in polylines:
+                xs, ys = kernels.scaled(p[0], xm.a, xm.b), kernels.scaled(p[1], ym.a, ym.b)
+                n = len(xs)
+                if n < 3:
+                    continue
+                seg = [math.hypot(xs[k + 1] - xs[k], ys[k + 1] - ys[k]) for k in range(n - 1)]
+                total = sum(seg)
+                if total < 3 * w:
+                    continue
+                # the label goes at the middle of the line
+                half, acc, k = total / 2, 0.0, 0
+                while k < n - 2 and acc + seg[k] < half:
+                    acc += seg[k]
+                    k += 1
+                t = (half - acc) / seg[k] if seg[k] else 0.0
+                cx = xs[k] + t * (xs[k + 1] - xs[k])
+                cy = ys[k] + t * (ys[k + 1] - ys[k])
+                if not (x0 <= cx <= x1 and y0 <= cy <= y1):
+                    continue
+                ang = math.degrees(math.atan2(-(ys[k + 1] - ys[k]), xs[k + 1] - xs[k]))
+                if ang > 90:
+                    ang -= 180
+                elif ang < -90:
+                    ang += 180
+                box = None if self._filled() else (w, a + d, 0.25 * size)
+                out.append((text, cx, cy, c, "center", "center", ang, box, None))
+                quads.append(quad(canvas, text, font, cx, cy, "center", "center", ang))
+        return [t for t, k in zip(out, drop(quads)) if k]
+
+
+@role("contour.bands")
+def _bands_paint(s: Contour, ctx: Any, glyph: bool, opacity: float = 1.0) -> dict[str, Any]:
+    """Filled bands in the colors of their middle values."""
+    axis = color_axis_of(s)
+    cmap, t = axis.color_map(), unit_of(axis)
+    return {"fill": [with_alpha(cmap(t((lo + hi) / 2)), opacity)
+                     for lo, hi, _ in s.stats["bands"]],
+            "simplify": s.resolve("simplify")}
+
+
+@role("contour.line")
+def _line_paint(s: Contour, ctx: Any, glyph: bool, i: int = 0, level: float = 0.0,
+                dash: Any = None, opacity: Any = None) -> dict[str, Any]:
+    """The lines of level ``i``."""
+    lw = s.resolve("line_width")
+    return {"color": with_alpha(s._level_color(ctx, i, level), opacity), "width": lw,
+            "dash": dash, "simplify": s.resolve("simplify")}
+
+
+@role("contour.key")
+def _key_paint(s: Contour, ctx: Any, glyph: bool) -> dict[str, Any]:
+    """The legend symbol: a swatch of the middle level's color (filled), or its line."""
+    c = s._key_color(ctx)
+    return {"fill": c, "color": c, "width": s.resolve("line_width")}

plot_io/presets/curves.py

@@ -0,0 +1,30 @@
+"""Line and fill interpolation and missing values (vec_io draws every mode in C:
+``interpolation=``, ``missing=`` on ``draw_lines`` and ``draw_area``)."""
+
+from typing import Any
+
+from ..settings import Setting, one_of
+
+__all__ = ["INTERPOLATIONS", "curve_settings"]
+
+INTERPOLATIONS = ("linear", "step-pre", "step-post", "step-mid", "monotone", "spline")
+
+
+def _interpolation(v: Any) -> str:
+    s = str(getattr(v, "value", v)).strip().lower().replace("_", "-")
+    if s not in INTERPOLATIONS:
+        raise ValueError(f"expected one of {', '.join(map(repr, INTERPOLATIONS))}, got {v!r}")
+    return s
+
+
+def curve_settings() -> dict[str, Setting]:
+    return {
+        "interpolation": Setting(
+            _interpolation,
+            doc="How points are joined: `linear`, `step-pre`, `step-post`, `step-mid`, "
+                "`monotone` (a curve that never overshoots) or `spline` (centripetal "
+                "Catmull–Rom)."),
+        "missing": Setting(one_of("break", "skip"),
+                           doc="Missing values: `break` the line there, or `skip` them and "
+                               "connect their neighbours."),
+    }