Coverage for src/pyhrp/dendrogram.py: 100%
221 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-10-14 19:39 +0000
1"""Hierarchical clustering tree construction and the Dendrogram container.
3This module builds the hierarchical clustering tree consumed by the HRP
4allocation entry points and stores it in a :class:`Dendrogram`:
5- build_tree: Build a hierarchical cluster tree from a correlation matrix
6- build_tree_from_operator: Build the tree from a covariance operator, matrix-free for single linkage
7- Dendrogram: Container for the clustering result and its visualization
8"""
10from __future__ import annotations
12from collections.abc import Generator, Sequence
13from dataclasses import dataclass
14from typing import TYPE_CHECKING, Literal
16import numpy as np
17import polars as pl
18import scipy.cluster.hierarchy as sch
19import scipy.spatial.distance as ssd
21from .algos import one_over_n as _one_over_n
22from .cluster import Cluster, Portfolio
23from .operators import CovarianceOperator, as_operator, block
25if TYPE_CHECKING:
26 import plotly.graph_objects as go
28__all__ = ["Dendrogram", "build_tree", "build_tree_from_operator"]
31@dataclass(frozen=True)
32class Dendrogram:
33 """Container for hierarchical clustering dendrogram data and visualization.
35 This class stores the results of hierarchical clustering and provides methods
36 for accessing and visualizing the dendrogram structure.
38 Attributes:
39 root (Cluster): The root node of the hierarchical clustering tree
40 assets (list[str]): Names of assets included in the clustering
41 linkage (np.ndarray | None): Linkage matrix in scipy format for plotting
42 distance (pl.DataFrame | None): Distance matrix used for clustering
43 method (str | None): Linkage method used for clustering
44 """
46 root: Cluster
47 assets: list[str]
48 distance: pl.DataFrame | None = None
49 linkage: np.ndarray | None = None
50 method: str | None = None
52 def __post_init__(self) -> None:
53 """Validate dataclass fields after initialization.
55 Ensures that the optional distance matrix, when provided, is a polars
56 DataFrame with columns aligned to the asset list, and verifies that the
57 number of leaves in the cluster tree matches the number of assets.
58 """
59 if self.distance is not None:
60 if not isinstance(self.distance, pl.DataFrame):
61 msg = "distance must be a polars DataFrame."
62 raise TypeError(msg)
64 if self.distance.columns != list(self.assets):
65 msg = "Distance matrix index/columns must align with assets."
66 raise ValueError(msg)
68 if len(self.root.leaves) != len(self.assets):
69 msg = "Number of leaves does not match number of assets."
70 raise ValueError(msg)
72 def plot(self, **kwargs: object) -> go.Figure:
73 """Build and return a plotly dendrogram figure.
75 Delegates to :func:`pyhrp.plot.plot_dendrogram`; the plotly dependency
76 is imported lazily so importing the allocation core stays plotly-free.
77 """
78 from .plot import plot_dendrogram
80 return plot_dendrogram(self, **kwargs)
82 def one_over_n(self) -> Generator[tuple[int, Portfolio]]:
83 """Yield the hierarchical 1/N portfolios level by level for this tree.
85 Container-level convenience wrapper around :func:`pyhrp.algos.one_over_n`.
87 Yields:
88 tuple[int, Portfolio]: The level number and the equal-weight portfolio
89 at that level.
90 """
91 yield from _one_over_n(self.root, self.assets)
93 @property
94 def ids(self) -> list[int]:
95 """Node values in the order left -> right as they appear in the dendrogram."""
96 return [node.value for node in self.root.leaves]
98 @property
99 def names(self) -> list[str]:
100 """The asset names as induced by the order of ids."""
101 return [self.assets[i] for i in self.ids]
104def _compute_distance_matrix(corr: pl.DataFrame) -> pl.DataFrame:
105 """Convert correlation matrix to distance matrix."""
106 c = corr.to_numpy()
107 dist = np.sqrt(np.clip((1.0 - c) / 2.0, a_min=0.0, a_max=1.0))
108 np.fill_diagonal(dist, 0.0)
109 cols = corr.columns
110 return pl.DataFrame(dict(zip(cols, dist, strict=True)))
113def _bisect_tree(ids: list[int], next_id: int) -> tuple[Cluster, int]:
114 """Build tree by recursive bisection."""
115 if not ids:
116 msg = "ids must contain at least one node id."
117 raise ValueError(msg)
118 if len(ids) == 1:
119 return Cluster(value=ids[0]), next_id
121 mid = len(ids) // 2
122 left_ids, right_ids = ids[:mid], ids[mid:]
123 left, next_id = _bisect_tree(left_ids, next_id)
124 right, next_id = _bisect_tree(right_ids, next_id)
125 next_id += 1
126 return Cluster(value=next_id, left=left, right=right), next_id
129def _cluster_diameter(node: Cluster, dist: np.ndarray) -> float:
130 """Largest pairwise distance between any two leaves under ``node``.
132 This is the merge height used for the bisection tree's linkage rows. The
133 diameter is monotone along the tree — a parent's leaf set is a superset of
134 each child's, so its maximum can only grow — which is what scipy's linkage
135 format requires of its distance column.
137 Args:
138 node (Cluster): The subtree whose leaves span the cluster.
139 dist (np.ndarray): Square distance matrix indexed by leaf value.
141 Returns:
142 float: The cluster diameter; 0.0 for a single leaf.
143 """
144 idx = [int(leaf.value) for leaf in node.leaves]
145 if len(idx) < 2:
146 return 0.0
147 return float(dist[np.ix_(idx, idx)].max())
150def _get_linkage(node: Cluster, dist: np.ndarray) -> list[list[float]]:
151 """Convert tree structure back to linkage matrix format.
153 Rows are emitted in post-order, matching the id assignment in
154 :func:`_bisect_tree`, so row ``i`` defines the cluster with id ``n + i``.
155 Column 2 carries the cluster diameter (see :func:`_cluster_diameter`),
156 because scipy reads that column as a cophenetic merge distance.
158 Args:
159 node (Cluster): The subtree to linearise.
160 dist (np.ndarray): Square distance matrix indexed by leaf value.
162 Returns:
163 list[list[float]]: Rows of ``[left id, right id, height, leaf count]``.
164 """
165 links_list: list[list[float]] = []
166 if node.left is not None and node.right is not None:
167 left, right = node._child_clusters()
168 links_list.extend(_get_linkage(left, dist))
169 links_list.extend(_get_linkage(right, dist))
170 links_list.append(
171 [
172 float(left.value),
173 float(right.value),
174 _cluster_diameter(node, dist),
175 float(len(left.leaves) + len(right.leaves)),
176 ]
177 )
178 return links_list
181def _check_finite_correlations(cor: pl.DataFrame, c: np.ndarray) -> None:
182 """Raise if the correlation matrix contains non-finite values.
184 Names the offending assets when the non-finite values sit on the diagonal,
185 since a constant (zero-variance) price series is the usual cause.
186 """
187 bad = [col for col, diag in zip(cor.columns, np.diagonal(c), strict=True) if not np.isfinite(diag)]
188 if bad:
189 msg = (
190 f"Correlation matrix contains non-finite values for assets {bad}; "
191 "constant (zero-variance) price series produce NaN correlations."
192 )
193 raise ValueError(msg)
194 if not np.isfinite(c).all():
195 msg = "Correlation matrix contains non-finite values."
196 raise ValueError(msg)
199def _validate_correlation_matrix(cor: pl.DataFrame) -> None:
200 """Validate the correlation matrix accepted by :func:`build_tree`.
202 Raises:
203 TypeError: If ``cor`` is not a polars DataFrame.
204 ValueError: If it has fewer than two assets or contains non-finite values.
205 """
206 if not isinstance(cor, pl.DataFrame):
207 msg = "Correlation matrix must be a polars DataFrame."
208 raise TypeError(msg)
209 if len(cor.columns) < 2:
210 msg = "Correlation matrix must contain at least two assets."
211 raise ValueError(msg)
212 _check_finite_correlations(cor, cor.to_numpy())
215def _to_cluster(node: sch.ClusterNode) -> Cluster:
216 """Convert a scipy ClusterNode tree into our Cluster format.
218 Args:
219 node (sch.ClusterNode): A node from scipy's hierarchical clustering.
221 Returns:
222 Cluster: Equivalent node in our Cluster format.
223 """
224 # Two passes instead of recursion. scipy's tree is as deep as the linkage is
225 # unbalanced -- single linkage chains, so depth grows with the asset count --
226 # and recursing here raised RecursionError before allocation was ever reached.
227 #
228 # Pass one collects nodes parent-before-child; reversing that order visits every
229 # child before its parent, which is what building bottom-up requires.
230 order: list[sch.ClusterNode] = []
231 stack: list[sch.ClusterNode] = [node]
232 while stack:
233 current = stack.pop()
234 order.append(current)
235 if current.left is not None and current.right is not None:
236 stack.append(current.left)
237 stack.append(current.right)
239 # Pass two builds each Cluster once its children exist. scipy assigns every
240 # node a unique id, so the id is a safe key for the partially built tree.
241 converted: dict[int, Cluster] = {}
242 for current in reversed(order):
243 if current.left is not None and current.right is not None:
244 converted[current.id] = Cluster(
245 value=current.id,
246 left=converted[current.left.id],
247 right=converted[current.right.id],
248 )
249 else:
250 converted[current.id] = Cluster(value=current.id)
252 return converted[node.id]
255def build_tree(
256 cor: pl.DataFrame, method: Literal["single", "complete", "average", "ward"] = "ward", bisection: bool = False
257) -> Dendrogram:
258 """Build hierarchical cluster tree from correlation matrix.
260 This function converts a correlation matrix to a distance matrix, performs
261 hierarchical clustering, and returns a Dendrogram object containing the
262 resulting tree structure.
264 Note:
265 There is no limit on the depth of the resulting tree. That matters because
266 ``method="single"`` chains: on smoothly decaying correlations it produces a
267 tree of depth *n* rather than *O(log n)*, so an ordinary equity universe can
268 exceed Python's recursion limit. Tree construction and every traversal that
269 walks the result are iterative for exactly that reason.
271 Args:
272 cor (pl.DataFrame): Correlation matrix of asset returns (columns are assets)
273 method (Literal["single", "complete", "average", "ward"]): Linkage method for hierarchical clustering
274 - "single": minimum distance between points (nearest neighbor)
275 - "complete": maximum distance between points (furthest neighbor)
276 - "average": average distance between all points
277 - "ward": Ward variance minimization
278 bisection (bool): Whether to use bisection method for tree construction
280 Returns:
281 Dendrogram: Object containing the hierarchical clustering tree, with:
282 - root: Root cluster node
283 - linkage: Linkage matrix for plotting
284 - assets: List of assets
285 - method: Clustering method used
286 - distance: Distance matrix
288 Examples:
289 >>> import polars as pl
290 >>> from pyhrp.dendrogram import build_tree
291 >>> cor = pl.DataFrame({"A": [1.0, 0.5], "B": [0.5, 1.0]})
292 >>> dg = build_tree(cor, method="ward")
293 >>> dg.root.leaf_count
294 2
295 """
296 _validate_correlation_matrix(cor)
297 dist = _compute_distance_matrix(cor)
298 links = sch.linkage(ssd.squareform(dist.to_numpy(), checks=False), method=method)
300 root = _to_cluster(sch.to_tree(links, rd=False))
302 # Apply bisection if requested
303 if bisection:
304 # Rebuild tree using bisection
305 leaf_ids: list[int] = [int(node.value) for node in root.leaves]
306 root, _ = _bisect_tree(ids=leaf_ids, next_id=max(leaf_ids))
307 links = np.array(_get_linkage(root, dist.to_numpy()))
309 return Dendrogram(root=root, linkage=links, method=method, distance=dist, assets=cor.columns)
312def _volatilities(op: CovarianceOperator) -> np.ndarray:
313 """Return ``sqrt(diag(Sigma))``, refusing an asset without positive variance."""
314 var = np.asarray(op.diag, dtype=np.float64)
315 if not np.all(var > 0.0):
316 bad = np.flatnonzero(~(var > 0.0)).tolist()
317 msg = f"Covariance has non-positive or non-finite variance at positions {bad}; correlations are undefined."
318 raise ValueError(msg)
319 return np.sqrt(var)
322def _unit(n: int, i: int) -> np.ndarray:
323 """The ``i``-th unit vector of length ``n``."""
324 e = np.zeros(n)
325 e[i] = 1.0
326 return e
329def _linkage_from_merges(merges: list[tuple[float, int, int]], n: int) -> np.ndarray:
330 """Turn merges into SciPy's linkage format.
332 Each merge ``(height, i, j)`` joins the clusters holding assets ``i`` and
333 ``j``. Merges are sorted by height and replayed by union-find, and each row
334 lists its two cluster ids smaller first, as SciPy labels them.
336 Args:
337 merges (list[tuple[float, int, int]]): The ``n - 1`` merges, in any order.
338 n (int): The number of assets.
340 Returns:
341 np.ndarray: The ``(n - 1) x 4`` linkage matrix.
342 """
343 # Union-find over the assets; label[r] is the cluster id of the set rooted at r.
344 root = np.arange(n)
345 label = np.arange(n)
346 size = np.ones(n, dtype=np.intp)
348 def find(i: int) -> int:
349 """Return the representative of ``i``, halving the path on the way."""
350 while root[i] != i:
351 root[i] = root[root[i]]
352 i = int(root[i])
353 return i
355 links = np.empty((n - 1, 4))
356 for k, (height, i, j) in enumerate(sorted(merges, key=lambda merge: merge[0])):
357 a, b = find(i), find(j)
358 links[k] = (min(label[a], label[b]), max(label[a], label[b]), height, size[a] + size[b])
359 root[b] = a
360 label[a] = n + k
361 size[a] += size[b]
362 return links
365def _single_linkage(op: CovarianceOperator, vol: np.ndarray) -> np.ndarray:
366 """Single linkage on ``d_ij = sqrt((1 - rho_ij) / 2)`` without the distance matrix.
368 Single linkage is the minimum spanning tree of the distance graph. Prim's
369 algorithm grows it from one column of ``Sigma`` per step, so memory is
370 ``O(n)`` beyond the operator and the time is ``n`` column products
371 (``O(n k)`` each for a rank-``k`` factor or Gram operator).
373 Args:
374 op (CovarianceOperator): The covariance.
375 vol (np.ndarray): The volatilities ``sqrt(diag(Sigma))``.
377 Returns:
378 np.ndarray: The ``(n - 1) x 4`` linkage matrix in SciPy's format.
379 """
380 n = op.n
381 in_tree = np.zeros(n, dtype=bool)
382 best = np.full(n, np.inf)
383 parent = np.zeros(n, dtype=np.intp)
384 edges: list[tuple[float, int, int]] = []
385 current = 0
386 for _ in range(n - 1):
387 in_tree[current] = True
388 rest = np.flatnonzero(~in_tree)
389 column = np.asarray(op.matvec(_unit(n, current)))[rest]
390 rho = column / (vol[rest] * vol[current])
391 dist = np.sqrt(np.clip((1.0 - rho) / 2.0, 0.0, 1.0))
392 closer = dist < best[rest]
393 best[rest[closer]] = dist[closer]
394 parent[rest[closer]] = current
395 current = int(rest[np.argmin(best[rest])])
396 edges.append((float(best[current]), int(parent[current]), current))
397 return _linkage_from_merges(edges, n)
400def _ward_linkage(op: CovarianceOperator, vol: np.ndarray) -> np.ndarray:
401 """Ward linkage on ``d_ij = sqrt((1 - rho_ij) / 2)`` without the distance matrix.
403 These distances are Euclidean: they are ``|p_i - p_j|`` for points with inner
404 products ``p_i . p_j = rho_ij / 4``. The centroid of a cluster ``A`` is then
405 reached through the vector ``v_A`` holding ``1 / (|A| sigma_i)`` on its assets,
406 with ``c_A . c_B = v_A' Sigma v_B / 4``. One product ``Sigma v_A`` therefore gives
407 the centroid inner products of ``A`` with every cluster at once (a sum over each
408 cluster's assets), and with them the Ward distance
410 ``D(A, B)**2 = 2 |A| |B| / (|A| + |B|) * |c_A - c_B|**2``,
412 which is the distance SciPy's Lance-Williams update reports. The nearest-
413 neighbour chain algorithm needs ``O(n)`` such queries, so the time is ``O(n)``
414 operator products (``O(n k)`` each for a rank-``k`` factor or Gram operator)
415 and the memory is ``O(n)``.
417 Args:
418 op (CovarianceOperator): The covariance.
419 vol (np.ndarray): The volatilities ``sqrt(diag(Sigma))``.
421 Returns:
422 np.ndarray: The ``(n - 1) x 4`` linkage matrix in SciPy's format.
423 """
424 n = op.n
425 # Every cluster is named by one of its assets, its representative; owner maps an
426 # asset to the representative of its cluster.
427 owner = np.arange(n)
428 size = np.ones(n)
429 norm2 = np.full(n, 0.25) # |c_A|**2; a single asset has rho_ii / 4
430 active = np.ones(n, dtype=bool)
431 merges: list[tuple[float, int, int]] = []
432 chain: list[int] = []
433 while len(merges) < n - 1:
434 if not chain:
435 chain.append(int(np.flatnonzero(active)[0]))
436 a = chain[-1]
437 members = np.flatnonzero(owner == a)
438 v = np.zeros(n)
439 v[members] = 1.0 / (size[a] * vol[members])
440 g = np.asarray(op.matvec(v))
441 dot = np.bincount(owner, weights=g / vol, minlength=n) / (4.0 * size)
442 d2 = 2.0 * size[a] * size / (size[a] + size) * np.maximum(norm2[a] + norm2 - 2.0 * dot, 0.0)
443 d2[~active] = np.inf
444 d2[a] = np.inf
445 b = int(np.argmin(d2))
446 # Reciprocal nearest neighbours: a's nearest is the cluster below it on the chain.
447 if len(chain) > 1 and d2[chain[-2]] <= d2[b]:
448 b = chain[-2]
449 del chain[-2:]
450 sa, sb = size[a], size[b]
451 norm2[a] = (sa * sa * norm2[a] + 2.0 * sa * sb * dot[b] + sb * sb * norm2[b]) / (sa + sb) ** 2
452 size[a] = sa + sb
453 owner[owner == b] = a
454 active[b] = False
455 merges.append((float(np.sqrt(d2[b])), a, b))
456 else:
457 chain.append(b)
458 return _linkage_from_merges(merges, n)
461def _bisect_linkage(root: Cluster, links: np.ndarray) -> tuple[Cluster, np.ndarray]:
462 """Rebuild a tree by bisection of its leaf order, without distances.
464 The height of a bisection node is the largest merge height between
465 neighbouring leaves in its range, which for a monotone tree is the range's
466 cophenetic diameter. Neighbouring leaves ``i`` and ``i + 1`` of a dendrogram
467 meet at the node that splits between them, so these gaps come from the
468 original tree in one pass. A parent's range contains its children's, so no
469 node sits below its children.
471 Args:
472 root (Cluster): The clustering tree (single or Ward linkage).
473 links (np.ndarray): Its ``(n - 1) x 4`` linkage matrix.
475 Returns:
476 tuple[Cluster, np.ndarray]: The bisection tree and its linkage matrix.
477 """
478 n = links.shape[0] + 1
479 leaves, spans = root._layout()
480 gaps = np.empty(n - 1)
481 for node, _, mid, _ in spans:
482 gaps[mid - 1] = links[int(node.value) - n, 2]
484 ids = [int(leaf.value) for leaf in leaves]
485 bisected, _ = _bisect_tree(ids=ids, next_id=n - 1)
486 rows = np.empty((n - 1, 4))
487 for node, lo, _, hi in bisected._layout()[1]:
488 left, right = node._child_clusters()
489 rows[int(node.value) - n] = (left.value, right.value, gaps[lo : hi - 1].max(), hi - lo)
490 return bisected, rows
493def build_tree_from_operator(
494 cov: CovarianceOperator,
495 assets: Sequence[str] | None = None,
496 method: Literal["single", "complete", "average", "ward"] = "ward",
497 bisection: bool = False,
498) -> Dendrogram:
499 """Build the hierarchical cluster tree from a covariance operator.
501 The correlation distance is ``d_ij = sqrt((1 - rho_ij) / 2)``, as in
502 :func:`build_tree`. Ward (the default) and single linkage are built
503 **matrix-free**, from the operator's diagonal and ``O(n)`` products with it
504 (see :func:`_ward_linkage` and :func:`_single_linkage`), never forming the
505 ``n x n`` correlation or distance matrix, so a factor or Gram operator
506 clusters a universe whose covariance would not fit in memory. They give
507 SciPy's trees. The returned :class:`Dendrogram` carries no ``distance``
508 frame.
510 With ``bisection=True`` the leaf order is split at its midpoint, as in
511 Lopez de Prado's original algorithm, giving ``ceil(log2(n))`` levels; this
512 matters for single linkage, whose tree chains to nearly ``n`` levels.
513 Matrix-free, the merge height of a bisection node is the cophenetic
514 diameter of its range in the clustering tree, not the largest pairwise
515 distance :func:`build_tree` reports; the tree is the same, only the plotted
516 heights differ.
518 Complete and average linkage need all pairwise distances; the correlation
519 is then materialised from the operator and the call is delegated to
520 :func:`build_tree`, at ``O(n**2)`` memory.
522 Args:
523 cov (CovarianceOperator): The covariance as an operator.
524 assets (Sequence[str], optional): Asset names; defaults to ``"0", "1", ...``.
525 method (Literal["single", "complete", "average", "ward"]): Linkage method.
526 Defaults to ``"ward"``; ``"ward"`` and ``"single"`` are matrix-free.
527 bisection (bool): Whether to rebuild the tree by bisection of the leaf order.
529 Returns:
530 Dendrogram: The clustering tree over the operator's assets.
532 Raises:
533 ValueError: If there are fewer than two assets, an asset has no positive
534 variance, or ``assets`` has the wrong length.
536 Examples:
537 >>> import numpy as np
538 >>> from pyhrp.dendrogram import build_tree_from_operator
539 >>> from pyhrp.operators import DenseCovariance
540 >>> sigma = np.array([[1.0, 0.9, 0.0], [0.9, 1.0, 0.0], [0.0, 0.0, 1.0]])
541 >>> dg = build_tree_from_operator(DenseCovariance(sigma), assets=["A", "B", "C"], method="single")
542 >>> dg.names
543 ['C', 'A', 'B']
544 """
545 op, names = as_operator(cov, assets)
546 if op.n < 2:
547 msg = "Covariance must contain at least two assets."
548 raise ValueError(msg)
549 vol = _volatilities(op)
551 if method in ("single", "ward"):
552 links = _single_linkage(op, vol) if method == "single" else _ward_linkage(op, vol)
553 root = _to_cluster(sch.to_tree(links, rd=False))
554 if bisection:
555 root, links = _bisect_linkage(root, links)
556 return Dendrogram(root=root, linkage=links, method=method, assets=names)
558 idx = np.arange(op.n)
559 corr = block(op, idx, idx) / np.outer(vol, vol)
560 np.fill_diagonal(corr, 1.0)
561 return build_tree(pl.DataFrame(dict(zip(names, corr, strict=True))), method=method, bisection=bisection)