plot-io 0.12.0__py3-none-win_arm64.whl

plot_io/helpers/tree.py

@@ -0,0 +1,214 @@
+"""Hierarchical clustering: :func:`linkage` (scipy's algorithms, in plot_io's C library) and
+:class:`Tree`, its merges with the leaf order and flat clusters."""
+
+import array
+from typing import Any
+
+from .. import kernels
+
+__all__ = ["Tree", "linkage"]
+
+
+class Tree:
+    """A binary tree of ``n`` observations as ``n - 1`` merges in scipy's layout: row ``i``
+    joins clusters ``left`` and ``right`` (observations ``0 … n - 1``, merges ``n + j``) at
+    ``height`` into a cluster of ``size`` observations.  ``merges`` is a scipy ``Z`` or any
+    table of such rows; ``labels`` names the observations.
+
+    ``merges`` is ``n − 1`` rows ``(left, right, height, size)`` as ``array('d')``; ``n``;
+    ``labels`` (``None`` when not given); ``leaves``, the observations in drawing order
+    (scipy's ``leaves_list``); ``cut(k=None, height=None)``, the flat cluster of each
+    observation, numbered from 1, as scipy's ``fcluster`` with ``maxclust`` (at most ``k``
+    clusters) or ``distance`` (merges up to ``height``).
+    """
+
+    __slots__ = ("merges", "n", "labels", "_leaves")
+
+    def __init__(self, merges: Any, labels: Any = None) -> None:
+        g = kernels.grid(merges)
+        if g is not None:
+            if g[1] != 4 or g[2] != 1:
+                raise ValueError("Tree: merges must have 4 columns")
+            rows = [g[3][4 * i:4 * i + 4] for i in range(g[0])]
+        else:
+            rows = [kernels.doubles(r) for r in merges]
+        n = len(rows) + 1
+        for i, r in enumerate(rows):
+            if len(r) != 4:
+                raise ValueError("Tree: merges must have 4 columns")
+            a, b = r[0], r[1]
+            if not (a == int(a) and b == int(b) and 0 <= a < n + i and 0 <= b < n + i
+                    and a != b):
+                raise ValueError(f"Tree: merge {i} joins {a!r} and {b!r}, not two earlier "
+                                 "clusters")
+        if labels is not None:
+            labels = list(labels)
+            if len(labels) != n:
+                raise ValueError(f"Tree: {len(labels)} labels for {n} observations")
+        #: the merges: rows of ``(left, right, height, size)`` (``array('d')``)
+        self.merges = rows
+        #: the number of observations
+        self.n = n
+        #: the names of the observations (``None``: none given)
+        self.labels = labels
+        self._leaves: list[int] | None = None
+
+    def __repr__(self) -> str:
+        return f"<Tree of {self.n} observations>"
+
+    @property
+    def leaves(self) -> list[int]:
+        """The observations in the order a dendrogram draws them (scipy's
+        ``leaves_list``: each merge's left cluster first)."""
+        if self._leaves is None:
+            n, out = self.n, []
+            stack = [2 * n - 2] if n > 1 else [0]
+            while stack:
+                k = stack.pop()
+                if k < n:
+                    out.append(k)
+                else:
+                    r = self.merges[k - n]
+                    stack += [int(r[1]), int(r[0])]
+            self._leaves = out
+        return list(self._leaves)
+
+    def _max_heights(self) -> list[float]:
+        """The largest height in each merge's subtree."""
+        n = self.n
+        out: list[float] = []
+        for r in self.merges:
+            h = r[2]
+            for c in (int(r[0]), int(r[1])):
+                if c >= n and out[c - n] > h:
+                    h = out[c - n]
+            out.append(h)
+        return out
+
+    def cut(self, k: int | None = None, height: float | None = None) -> list[int]:
+        """Flat clusters, numbered from 1, of each observation (scipy's ``fcluster``): at
+        most ``k`` clusters (``criterion="maxclust"``), or the subtrees whose merges are all
+        at most ``height`` (``criterion="distance"``)."""
+        if (k is None) == (height is None):
+            raise ValueError("cut: give k or height")
+        n = self.n
+        if n < 2:
+            return [1] * n
+        md = self._max_heights()
+        if k is not None:
+            if int(k) != k or k < 1:
+                raise ValueError(f"cut: k must be a positive integer, not {k!r}")
+            ts = [float("-inf")] + sorted(md)     # the count falls as the threshold rises
+            lo, hi = 0, len(ts) - 1
+            while lo < hi:
+                mid = (lo + hi) // 2
+                if self._count(md, ts[mid]) <= k:
+                    hi = mid
+                else:
+                    lo = mid + 1
+            return self._flat(md, ts[lo])
+        return self._flat(md, float(height))           # type: ignore[arg-type]
+
+    def _count(self, md: list[float], t: float) -> int:
+        """The number of flat clusters at threshold ``t``."""
+        n = self.n
+        count, stack = 0, [2 * n - 2]
+        while stack:
+            c = stack.pop()
+            if c < n:
+                count += 1
+            elif md[c - n] <= t:
+                count += 1
+            else:
+                r = self.merges[c - n]
+                stack += [int(r[0]), int(r[1])]
+        return count
+
+    def _flat(self, md: list[float], t: float) -> list[int]:
+        """scipy's ``cluster_monocrit``: the labels in its traversal order."""
+        n, merges = self.n, self.merges
+        out = [0] * n
+        visited: set[int] = set()
+        nodes, leader, count = [2 * n - 2], -1, 0
+        while nodes:
+            root = nodes[-1] - n
+            lc, rc = int(merges[root][0]), int(merges[root][1])
+            if leader == -1 and md[root] <= t:
+                leader = root
+                count += 1
+            if lc >= n and lc not in visited:
+                visited.add(lc)
+                nodes.append(lc)
+                continue
+            if rc >= n and rc not in visited:
+                visited.add(rc)
+                nodes.append(rc)
+                continue
+            for c in (lc, rc):
+                if c < n:
+                    if leader == -1:
+                        count += 1
+                    out[c] = count
+            if leader == root:
+                leader = -1
+            nodes.pop()
+        return out
+
+
+def _frame(data: Any) -> tuple[Any, Any, Any]:
+    """A data frame's values with its row and column names (pandas, polars), else
+    ``data``."""
+    if hasattr(data, "to_numpy") and hasattr(data, "columns"):
+        index = getattr(data, "index", None)
+        rows = list(index) if index is not None and hasattr(index, "__iter__") else None
+        return data.to_numpy(dtype=float), rows, list(data.columns)
+    return data, None, None
+
+
+def linkage(data: Any, method: str = "average", metric: str = "euclidean",
+            axis: str = "rows", labels: Any = None) -> Tree:
+    """The :class:`Tree` of a hierarchical clustering of the rows (``axis="cols"``: the
+    columns) of ``data``, a 2-D buffer, a sequence of rows or a data frame (its row or
+    column names the labels); ``metric="precomputed"``: ``data`` is the condensed distance
+    vector.  Methods: ``single``, ``complete``, ``average``, ``weighted``, ``ward``,
+    ``centroid``, ``median``; metrics: ``euclidean``, ``sqeuclidean``, ``cityblock``,
+    ``chebyshev``, ``cosine``, ``correlation``, ``precomputed`` (ward, centroid and median:
+    euclidean or precomputed).  The merges equal scipy's ``linkage`` on data without equal
+    distances; equal distances merge in a fixed order.  At most 16,384 observations.
+
+    The algorithms are scipy's, in plot_io's C library; ``data`` with ``precomputed`` is
+    the condensed distance vector as scipy's ``pdist`` gives.  Equal distances merge in a
+    fixed order, the same on every run and platform but not always scipy's.  NaN or
+    infinite values, a constant row with ``correlation``, a zero row with ``cosine`` and
+    more than 16,384 observations raise.  The :class:`Tree` gives ``merges``, ``n``,
+    ``labels``, ``leaves`` and ``cut``; :class:`Dendrogram` draws it.
+
+    Example::
+
+        rng = random.Random(1)
+        rows = [[rng.gauss(c, 1) for _ in range(4)] for c in (0, 0, 0, 4, 4, 8)]
+        tree = plt.linkage(rows, method="average")
+        clusters = tree.cut(k=3)
+    """
+    if axis not in ("rows", "cols"):
+        raise ValueError(f"linkage: axis is 'rows' or 'cols', not {axis!r}")
+    values, row_names, col_names = _frame(data)
+    if labels is None:
+        labels = col_names if axis == "cols" else row_names
+    if axis == "cols" and metric != "precomputed":
+        g = kernels.grid(values)
+        if g is not None:
+            h, w, flat = g[0], g[1], g[3]
+        else:
+            rs = [kernels.doubles(r) for r in values]
+            h, w = len(rs), len(rs[0]) if rs else 0
+            flat = array.array("d")
+            for r in rs:
+                if len(r) != w:
+                    raise ValueError("linkage: rows of different lengths")
+                flat.extend(r)
+        values = [flat[j::w] for j in range(w)] if h else []
+    n, z = kernels.linkage(values, method, metric)
+    if n < 1:
+        raise ValueError("linkage: no observations")
+    return Tree([z[4 * i:4 * i + 4] for i in range(n - 1)], labels)