Coverage for src/pyhrp/dendrogram.py: 100%

1"""Hierarchical clustering tree construction and the Dendrogram container. 

2 

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- Dendrogram: Container for the clustering result and its visualization 

7""" 

8 

9from __future__ import annotations 

10 

11from collections.abc import Generator 

12from dataclasses import dataclass 

13from typing import TYPE_CHECKING, Literal 

14 

15import numpy as np 

16import polars as pl 

17import scipy.cluster.hierarchy as sch 

18import scipy.spatial.distance as ssd 

19 

20from .algos import one_over_n as _one_over_n 

21from .cluster import Cluster, Portfolio 

22 

23if TYPE_CHECKING: 

24 import plotly.graph_objects as go 

25 

26__all__ = ["Dendrogram", "build_tree"] 

27 

28 

29@dataclass(frozen=True) 

30class Dendrogram: 

31 """Container for hierarchical clustering dendrogram data and visualization. 

32 

33 This class stores the results of hierarchical clustering and provides methods 

34 for accessing and visualizing the dendrogram structure. 

35 

36 Attributes: 

37 root (Cluster): The root node of the hierarchical clustering tree 

38 assets (list[str]): Names of assets included in the clustering 

39 linkage (np.ndarray | None): Linkage matrix in scipy format for plotting 

40 distance (pl.DataFrame | None): Distance matrix used for clustering 

41 method (str | None): Linkage method used for clustering 

42 """ 

43 

44 root: Cluster 

45 assets: list[str] 

46 distance: pl.DataFrame | None = None 

47 linkage: np.ndarray | None = None 

48 method: str | None = None 

49 

50 def __post_init__(self) -> None: 

51 """Validate dataclass fields after initialization. 

52 

53 Ensures that the optional distance matrix, when provided, is a polars 

54 DataFrame with columns aligned to the asset list, and verifies that the 

55 number of leaves in the cluster tree matches the number of assets. 

56 """ 

57 if self.distance is not None: 

58 if not isinstance(self.distance, pl.DataFrame): 

59 msg = "distance must be a polars DataFrame." 

60 raise TypeError(msg) 

61 

62 if self.distance.columns != list(self.assets): 

63 msg = "Distance matrix index/columns must align with assets." 

64 raise ValueError(msg) 

65 

66 if len(self.root.leaves) != len(self.assets): 

67 msg = "Number of leaves does not match number of assets." 

68 raise ValueError(msg) 

69 

70 def plot(self, **kwargs: object) -> go.Figure: 

71 """Build and return a plotly dendrogram figure. 

72 

73 Delegates to :func:`pyhrp.plot.plot_dendrogram`; the plotly dependency 

74 is imported lazily so importing the allocation core stays plotly-free. 

75 """ 

76 from .plot import plot_dendrogram 

77 

78 return plot_dendrogram(self, **kwargs) 

79 

80 def one_over_n(self) -> Generator[tuple[int, Portfolio]]: 

81 """Yield the hierarchical 1/N portfolios level by level for this tree. 

82 

83 Container-level convenience wrapper around :func:`pyhrp.algos.one_over_n`. 

84 

85 Yields: 

86 tuple[int, Portfolio]: The level number and the equal-weight portfolio 

87 at that level. 

88 """ 

89 yield from _one_over_n(self.root, self.assets) 

90 

91 @property 

92 def ids(self) -> list[int]: 

93 """Node values in the order left -> right as they appear in the dendrogram.""" 

94 return [node.value for node in self.root.leaves] 

95 

96 @property 

97 def names(self) -> list[str]: 

98 """The asset names as induced by the order of ids.""" 

99 return [self.assets[i] for i in self.ids] 

100 

101 

102def _compute_distance_matrix(corr: pl.DataFrame) -> pl.DataFrame: 

103 """Convert correlation matrix to distance matrix.""" 

104 c = corr.to_numpy() 

105 dist = np.sqrt(np.clip((1.0 - c) / 2.0, a_min=0.0, a_max=1.0)) 

106 np.fill_diagonal(dist, 0.0) 

107 cols = corr.columns 

108 return pl.DataFrame(dict(zip(cols, dist, strict=True))) 

109 

110 

111def _bisect_tree(ids: list[int], next_id: int) -> tuple[Cluster, int]: 

112 """Build tree by recursive bisection.""" 

113 if not ids: 

114 msg = "ids must contain at least one node id." 

115 raise ValueError(msg) 

116 if len(ids) == 1: 

117 return Cluster(value=ids[0]), next_id 

118 

119 mid = len(ids) // 2 

120 left_ids, right_ids = ids[:mid], ids[mid:] 

121 left, next_id = _bisect_tree(left_ids, next_id) 

122 right, next_id = _bisect_tree(right_ids, next_id) 

123 next_id += 1 

124 return Cluster(value=next_id, left=left, right=right), next_id 

125 

126 

127def _cluster_diameter(node: Cluster, dist: np.ndarray) -> float: 

128 """Largest pairwise distance between any two leaves under ``node``. 

129 

130 This is the merge height used for the bisection tree's linkage rows. The 

131 diameter is monotone along the tree — a parent's leaf set is a superset of 

132 each child's, so its maximum can only grow — which is what scipy's linkage 

133 format requires of its distance column. 

134 

135 Args: 

136 node (Cluster): The subtree whose leaves span the cluster. 

137 dist (np.ndarray): Square distance matrix indexed by leaf value. 

138 

139 Returns: 

140 float: The cluster diameter; 0.0 for a single leaf. 

141 """ 

142 idx = [int(leaf.value) for leaf in node.leaves] 

143 if len(idx) < 2: 

144 return 0.0 

145 return float(dist[np.ix_(idx, idx)].max()) 

146 

147 

148def _get_linkage(node: Cluster, dist: np.ndarray) -> list[list[float]]: 

149 """Convert tree structure back to linkage matrix format. 

150 

151 Rows are emitted in post-order, matching the id assignment in 

152 :func:`_bisect_tree`, so row ``i`` defines the cluster with id ``n + i``. 

153 Column 2 carries the cluster diameter (see :func:`_cluster_diameter`), 

154 because scipy reads that column as a cophenetic merge distance. 

155 

156 Args: 

157 node (Cluster): The subtree to linearise. 

158 dist (np.ndarray): Square distance matrix indexed by leaf value. 

159 

160 Returns: 

161 list[list[float]]: Rows of ``[left id, right id, height, leaf count]``. 

162 """ 

163 links_list: list[list[float]] = [] 

164 if node.left is not None and node.right is not None: 

165 left, right = node._child_clusters() 

166 links_list.extend(_get_linkage(left, dist)) 

167 links_list.extend(_get_linkage(right, dist)) 

168 links_list.append( 

169 [ 

170 float(left.value), 

171 float(right.value), 

172 _cluster_diameter(node, dist), 

173 float(len(left.leaves) + len(right.leaves)), 

174 ] 

175 ) 

176 return links_list 

177 

178 

179def _check_finite_correlations(cor: pl.DataFrame, c: np.ndarray) -> None: 

180 """Raise if the correlation matrix contains non-finite values. 

181 

182 Names the offending assets when the non-finite values sit on the diagonal, 

183 since a constant (zero-variance) price series is the usual cause. 

184 """ 

185 bad = [col for col, diag in zip(cor.columns, np.diagonal(c), strict=True) if not np.isfinite(diag)] 

186 if bad: 

187 msg = ( 

188 f"Correlation matrix contains non-finite values for assets {bad}; " 

189 "constant (zero-variance) price series produce NaN correlations." 

190 ) 

191 raise ValueError(msg) 

192 if not np.isfinite(c).all(): 

193 msg = "Correlation matrix contains non-finite values." 

194 raise ValueError(msg) 

195 

196 

197def _validate_correlation_matrix(cor: pl.DataFrame) -> None: 

198 """Validate the correlation matrix accepted by :func:`build_tree`. 

199 

200 Raises: 

201 TypeError: If ``cor`` is not a polars DataFrame. 

202 ValueError: If it has fewer than two assets or contains non-finite values. 

203 """ 

204 if not isinstance(cor, pl.DataFrame): 

205 msg = "Correlation matrix must be a polars DataFrame." 

206 raise TypeError(msg) 

207 if len(cor.columns) < 2: 

208 msg = "Correlation matrix must contain at least two assets." 

209 raise ValueError(msg) 

210 _check_finite_correlations(cor, cor.to_numpy()) 

211 

212 

213def _to_cluster(node: sch.ClusterNode) -> Cluster: 

214 """Convert a scipy ClusterNode tree into our Cluster format. 

215 

216 Args: 

217 node (sch.ClusterNode): A node from scipy's hierarchical clustering. 

218 

219 Returns: 

220 Cluster: Equivalent node in our Cluster format. 

221 """ 

222 # Two passes instead of recursion. scipy's tree is as deep as the linkage is 

223 # unbalanced -- single linkage chains, so depth grows with the asset count -- 

224 # and recursing here raised RecursionError before allocation was ever reached. 

225 # 

226 # Pass one collects nodes parent-before-child; reversing that order visits every 

227 # child before its parent, which is what building bottom-up requires. 

228 order: list[sch.ClusterNode] = [] 

229 stack: list[sch.ClusterNode] = [node] 

230 while stack: 

231 current = stack.pop() 

232 order.append(current) 

233 if current.left is not None and current.right is not None: 

234 stack.append(current.left) 

235 stack.append(current.right) 

236 

237 # Pass two builds each Cluster once its children exist. scipy assigns every 

238 # node a unique id, so the id is a safe key for the partially built tree. 

239 converted: dict[int, Cluster] = {} 

240 for current in reversed(order): 

241 if current.left is not None and current.right is not None: 

242 converted[current.id] = Cluster( 

243 value=current.id, 

244 left=converted[current.left.id], 

245 right=converted[current.right.id], 

246 ) 

247 else: 

248 converted[current.id] = Cluster(value=current.id) 

249 

250 return converted[node.id] 

251 

252 

253def build_tree( 

254 cor: pl.DataFrame, method: Literal["single", "complete", "average", "ward"] = "ward", bisection: bool = False 

255) -> Dendrogram: 

256 """Build hierarchical cluster tree from correlation matrix. 

257 

258 This function converts a correlation matrix to a distance matrix, performs 

259 hierarchical clustering, and returns a Dendrogram object containing the 

260 resulting tree structure. 

261 

262 Args: 

263 cor (pl.DataFrame): Correlation matrix of asset returns (columns are assets) 

264 method (Literal["single", "complete", "average", "ward"]): Linkage method for hierarchical clustering 

265 - "single": minimum distance between points (nearest neighbor) 

266 - "complete": maximum distance between points (furthest neighbor) 

267 - "average": average distance between all points 

268 - "ward": Ward variance minimization 

269 bisection (bool): Whether to use bisection method for tree construction 

270 

271 Returns: 

272 Dendrogram: Object containing the hierarchical clustering tree, with: 

273 - root: Root cluster node 

274 - linkage: Linkage matrix for plotting 

275 - assets: List of assets 

276 - method: Clustering method used 

277 - distance: Distance matrix 

278 

279 Examples: 

280 >>> import polars as pl 

281 >>> from pyhrp.dendrogram import build_tree 

282 >>> cor = pl.DataFrame({"A": [1.0, 0.5], "B": [0.5, 1.0]}) 

283 >>> dg = build_tree(cor, method="ward") 

284 >>> dg.root.leaf_count 

285 2 

286 """ 

287 _validate_correlation_matrix(cor) 

288 dist = _compute_distance_matrix(cor) 

289 links = sch.linkage(ssd.squareform(dist.to_numpy(), checks=False), method=method) 

290 

291 root = _to_cluster(sch.to_tree(links, rd=False)) 

292 

293 # Apply bisection if requested 

294 if bisection: 

295 # Rebuild tree using bisection 

296 leaf_ids: list[int] = [int(node.value) for node in root.leaves] 

297 root, _ = _bisect_tree(ids=leaf_ids, next_id=max(leaf_ids)) 

298 links = np.array(_get_linkage(root, dist.to_numpy())) 

299 

300 return Dendrogram(root=root, linkage=links, method=method, distance=dist, assets=cor.columns)