plot-io 0.12.0__py3-none-win_amd64.whl

plot_io/occupancy.py

@@ -0,0 +1,295 @@
+"""What the series cover inside the axes, on a coarse grid: used to place legends."""
+
+import math
+from typing import Any
+
+from . import kernels
+
+
+__all__ = ["Occupancy", "thin", "extremes", "N", "MAX_VERTICES", "MAX_POINTS"]
+
+#: Cells across each side of the axes.
+N = 64
+#: Polylines and curves are thinned to at most this many vertices before they are marked.
+MAX_VERTICES = 4096
+
+
+#: Points beyond this many are sampled (every k-th) before they are marked.
+MAX_POINTS = 50_000
+
+
+def _stride(n: int) -> int:
+    return max(1, -(-n // MAX_VERTICES))
+
+
+def thin(*columns: Any, limit: int = MAX_POINTS) -> tuple[Any, ...]:
+    """Every k-th item of each column, so that at most ``limit`` are left (data are thinned
+    before they are mapped to the page, which keeps legend placement fast)."""
+    n = min(len(c) for c in columns)
+    k = max(1, -(-n // limit))
+    if k == 1:
+        return columns
+    return tuple(c[:n:k] for c in columns)
+
+
+def extremes(a: Any, b: Any, limit: int = MAX_VERTICES) -> tuple[Any, Any]:
+    """A polyline of more than ``4 · limit`` vertices thinned to strides of ``n / limit``
+    vertices, each kept by its first vertex and those of least and largest ``a`` and ``b``
+    (in order, so that narrow spikes stay); a stride with NaN keeps its first NaN (a
+    break) and the extremes of its other vertices."""
+    n = min(len(a), len(b))
+    k = -(-n // limit)
+    if k <= 4:
+        return a[:n], b[:n]
+    keep: list[int] = []
+    for i in range(0, n, k):
+        sa, sb = a[i:i + k], b[i:i + k]
+        ta, tb = sum(sa), sum(sb)
+        if ta == ta and tb == tb and abs(ta) != math.inf and abs(tb) != math.inf:
+            idx = {0, sa.index(min(sa)), sa.index(max(sa)), sb.index(min(sb)),
+                   sb.index(max(sb))}
+        else:
+            good = [m for m in range(len(sa)) if math.isfinite(sa[m]) and math.isfinite(sb[m])]
+            idx = {m for m in range(len(sa)) if not (sa[m] == sa[m] and sb[m] == sb[m])}
+            idx = {min(idx)} if idx else set()
+            if good:
+                idx |= {good[0], min(good, key=sa.__getitem__), max(good, key=sa.__getitem__),
+                        min(good, key=sb.__getitem__), max(good, key=sb.__getitem__)}
+        keep.extend(i + m for m in sorted(idx))
+    return [a[i] for i in keep], [b[i] for i in keep]
+
+
+class Occupancy:
+    """An ``N`` × ``N`` grid of cells over the axes rectangle (page units); a cell is
+    occupied when some series draws in it.  ``marks``: the same grid of the cells under
+    markers and solid boxes only (texts and guides inside the axes; a marker or a text
+    under a legend is hidden, a line or an area is not); ``discs`` and ``solids``: those
+    markers (page positions and radius) and boxes themselves, which :meth:`solid_in` tests
+    a box against exactly."""
+
+    def __init__(self, rect: tuple[float, float, float, float]) -> None:
+        x0, y0, x1, y1 = rect
+        self.x0, self.y0 = x0, y0
+        self.cw = max(x1 - x0, 1e-9) / N
+        self.ch = max(y1 - y0, 1e-9) / N
+        self.cells = bytearray(N * N)
+        self.marks = bytearray(N * N)
+        self.discs: list[tuple[Any, Any, float]] = []
+        self.solids: list[tuple[float, float, float, float]] = []
+        self._xe = [x0 + i * self.cw for i in range(N + 1)]
+        self._ye = [y0 + j * self.ch for j in range(N + 1)]
+
+    # -- marking ---------------------------------------------------------------------
+
+    def _col(self, x: float) -> int:
+        return min(N - 1, max(0, int((x - self.x0) / self.cw)))
+
+    def _row(self, y: float) -> int:
+        return min(N - 1, max(0, int((y - self.y0) / self.ch)))
+
+    def _span(self, r: int, c0: int, c1: int, grid: bytearray | None = None) -> None:
+        if c1 >= c0:
+            (self.cells if grid is None else grid)[r * N + c0:r * N + c1 + 1] = \
+                b"\x01" * (c1 - c0 + 1)
+
+    def points(self, px: Any, py: Any, radius: float = 0.0, marker: bool = False) -> None:
+        """Marks the cells holding points (counted in C; points outside are ignored), and
+        the cells within ``radius`` (px, a marker's half size) of those; ``marker``: in
+        ``marks`` too."""
+        if not len(px):
+            return
+        if marker:
+            self.discs.append((px, py, radius))
+        _, _, counts = kernels.histogram2d(px, py, self._xe, self._ye)
+        kx = int(radius / self.cw + 0.5)
+        ky = int(radius / self.ch + 0.5)
+        grids = (self.cells, self.marks) if marker else (self.cells,)
+        for k, c in enumerate(counts):
+            if not c:
+                continue
+            if not kx and not ky:
+                for g in grids:
+                    g[k] = 1
+                continue
+            r, col = divmod(k, N)
+            c0, c1 = max(0, col - kx), min(N - 1, col + kx)
+            for rr in range(max(0, r - ky), min(N - 1, r + ky) + 1):
+                for g in grids:
+                    self._span(rr, c0, c1, g)
+
+    def lines(self, px: Any, py: Any) -> None:
+        """Marks the cells a polyline passes through (NaN breaks it; a long one thinned by
+        :func:`extremes`); segments are sampled at half a cell, so long segments count
+        along their whole length (at most ``8 · MAX_VERTICES`` samples in all, spread by
+        length)."""
+        n = min(len(px), len(py))
+        if n == 0:
+            return
+        verts = list(zip(*extremes(px, py)))
+        step = 0.5 * min(self.cw, self.ch)
+        counts = []
+        for (x0, y0), (x1, y1) in zip(verts, verts[1:]):
+            d = max(abs(x1 - x0), abs(y1 - y0))
+            counts.append(int(d / step) if d == d else -1)
+        total = sum(c for c in counts if c > 0)
+        budget = 8 * MAX_VERTICES
+        scale = budget / total if total > budget else 1.0
+        xs, ys = [], []
+        for (x0, y0), (x1, y1), k in zip(verts, verts[1:], counts):
+            if k < 0:
+                continue
+            k = int(k * scale)
+            for j in range(k + 1):
+                t = j / (k + 1)
+                xs.append(x0 + (x1 - x0) * t)
+                ys.append(y0 + (y1 - y0) * t)
+        x, y = verts[-1]
+        if x == x and y == y:
+            xs.append(x)
+            ys.append(y)
+        self.points(xs, ys)
+
+    def rects(self, x0: Any, y0: Any, x1: Any, y1: Any, solid: bool = False) -> None:
+        """Marks the cells under rectangles (any corner order; NaN rectangles skipped);
+        ``solid``: texts and guides, in ``marks`` and ``solids`` too."""
+        X0, Y0 = self.x0, self.y0
+        X1, Y1 = X0 + N * self.cw, Y0 + N * self.ch
+        for a, b, c, d in zip(x0, y0, x1, y1):
+            if a != a or b != b or c != c or d != d:
+                continue
+            a, c = min(a, c), max(a, c)
+            b, d = min(b, d), max(b, d)
+            if c < X0 or a > X1 or d < Y0 or b > Y1:
+                continue
+            c0, c1 = self._col(a), self._col(c)
+            for r in range(self._row(b), self._row(d) + 1):
+                self._span(r, c0, c1)
+                if solid:
+                    self._span(r, c0, c1, self.marks)
+            if solid:
+                self.solids.append((a, b, c, d))
+
+    def band(self, t: Any, lo: Any, hi: Any, horizontal: bool = False) -> None:
+        """Marks the region between two curves sharing the coordinate ``t`` (page x, or page
+        y when ``horizontal``): per grid column (row), from the least to the largest value
+        of the samples in it, with samples added at half a cell along ``t``."""
+        n = min(len(t), len(lo), len(hi))
+        if n == 0:
+            return
+        size = self.ch if horizontal else self.cw
+        index = self._row if horizontal else self._col
+        mins = [math.inf] * N
+        maxs = [-math.inf] * N
+        s = _stride(n)
+        prev: tuple[float, float, float] | None = None
+        T0 = self.y0 if horizontal else self.x0
+        T1 = T0 + N * size
+
+        def put(tv: float, a: float, b: float) -> None:
+            if not T0 <= tv <= T1:                  # outside the view
+                return
+            k = index(tv)
+            lo_, hi_ = (a, b) if a <= b else (b, a)
+            if lo_ < mins[k]:
+                mins[k] = lo_
+            if hi_ > maxs[k]:
+                maxs[k] = hi_
+
+        for i in range(0, n, s):
+            tv, a, b = t[i], lo[i], hi[i]
+            if tv != tv or a != a or b != b:
+                prev = None
+                continue
+            if prev is not None:
+                k = int(abs(tv - prev[0]) / (0.5 * size))
+                for j in range(1, k + 1):
+                    f = j / (k + 1)
+                    put(prev[0] + (tv - prev[0]) * f, prev[1] + (a - prev[1]) * f,
+                        prev[2] + (b - prev[2]) * f)
+            put(tv, a, b)
+            prev = (tv, a, b)
+        for k in range(N):
+            if mins[k] > maxs[k]:
+                continue
+            centre = T0 + (k + 0.5) * size
+            if not T0 <= centre <= T1:
+                continue
+            if horizontal:
+                self._span(k, self._col(mins[k]), self._col(maxs[k]))
+            else:
+                for r in range(self._row(mins[k]), self._row(maxs[k]) + 1):
+                    self.cells[r * N + k] = 1
+
+    # -- queries ---------------------------------------------------------------------
+
+    def _sums(self) -> list[int]:
+        """Summed-area table: ``s[(r) * (N + 1) + c]`` = occupied cells above and left."""
+        s = [0] * ((N + 1) * (N + 1))
+        cells = self.cells
+        for r in range(N):
+            run = 0
+            base, prev = (r + 1) * (N + 1), r * (N + 1)
+            for c in range(N):
+                run += cells[r * N + c]
+                s[base + c + 1] = s[prev + c + 1] + run
+        return s
+
+    def covered(self, x: float, y: float, w: float, h: float, sums: list[int] | None = None,
+                markers: bool = False) -> float:
+        """Fraction of the cells whose centres lie in the box ``(x, y, w, h)`` that are
+        occupied (``markers``: under markers; 0 when the box holds no cell centre).
+        ``sums`` (from :meth:`_sums`, of all cells) makes repeated queries constant time."""
+        c0 = max(0, math.ceil((x - self.x0) / self.cw - 0.5))
+        c1 = min(N - 1, math.floor((x + w - self.x0) / self.cw - 0.5))
+        r0 = max(0, math.ceil((y - self.y0) / self.ch - 0.5))
+        r1 = min(N - 1, math.floor((y + h - self.y0) / self.ch - 0.5))
+        if c1 < c0 or r1 < r0:
+            return 0.0
+        if sums is None or markers:
+            cells = self.marks if markers else self.cells
+            hit = sum(cells[r * N + c0:r * N + c1 + 1].count(1) for r in range(r0, r1 + 1))
+        else:
+            M = N + 1
+            hit = (sums[(r1 + 1) * M + c1 + 1] - sums[r0 * M + c1 + 1]
+                   - sums[(r1 + 1) * M + c0] + sums[r0 * M + c0])
+        return hit / ((c1 - c0 + 1) * (r1 - r0 + 1))
+
+    def solid_in(self, x: float, y: float, w: float, h: float) -> int:
+        """The markers whose disc meets the box ``(x, y, w, h)`` (by the box grown by
+        each marker's radius) and the solid boxes it meets."""
+        n = sum(1 for a, b, c, d in self.solids if a < x + w and x < c and b < y + h and y < d)
+        for px, py, r in self.discs:
+            r = max(r, 0.5)
+            n += int(kernels.histogram2d(px, py, (x - r, x + w + r),
+                                         (y - r, y + h + r))[2][0])
+        return n
+
+    def free_boxes(self, w: float, h: float, pad: float) -> list[tuple[float, float]]:
+        """Top-left corners, one cell apart, of the boxes ``w`` × ``h`` inside the axes
+        (``pad`` from their edges) that cover no occupied cell."""
+        sums = self._sums()
+        x_lo, x_hi = self.x0 + pad, self.x0 + N * self.cw - pad - w
+        y_lo, y_hi = self.y0 + pad, self.y0 + N * self.ch - pad - h
+        if x_hi < x_lo or y_hi < y_lo:
+            return []
+        nx = int((x_hi - x_lo) / self.cw)
+        ny = int((y_hi - y_lo) / self.ch)
+        xs = [x_lo + i * self.cw for i in range(nx + 1)] + [x_hi]
+        ys = [y_lo + j * self.ch for j in range(ny + 1)] + [y_hi]
+        # the cells whose centres a box at each x (y) holds, as :meth:`covered` takes them
+        cols = [(max(0, math.ceil((x - self.x0) / self.cw - 0.5)),
+                 min(N - 1, math.floor((x + w - self.x0) / self.cw - 0.5))) for x in xs]
+        rows = [(max(0, math.ceil((y - self.y0) / self.ch - 0.5)),
+                 min(N - 1, math.floor((y + h - self.y0) / self.ch - 0.5))) for y in ys]
+        M = N + 1
+        out = []
+        for y, (r0, r1) in zip(ys, rows):
+            if r1 < r0:
+                out += [(x, y) for x in xs]
+                continue
+            top, bottom = r0 * M, (r1 + 1) * M
+            for x, (c0, c1) in zip(xs, cols):
+                if c1 < c0 or (sums[bottom + c1 + 1] - sums[top + c1 + 1]
+                               - sums[bottom + c0] + sums[top + c0]) == 0:
+                    out.append((x, y))
+        return out