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 .cluster import Cluster, Portfolio 

21 

22if TYPE_CHECKING: 

23 import plotly.graph_objects as go 

24 

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

26 

27 

28@dataclass(frozen=True) 

29class Dendrogram: 

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

31 

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

33 for accessing and visualizing the dendrogram structure. 

34 

35 Attributes: 

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

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

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

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

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

41 """ 

42 

43 root: Cluster 

44 assets: list[str] 

45 distance: pl.DataFrame | None = None 

46 linkage: np.ndarray | None = None 

47 method: str | None = None 

48 

49 def __post_init__(self) -> None: 

50 """Validate dataclass fields after initialization. 

51 

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

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

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

55 """ 

56 if self.distance is not None: 

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

58 raise TypeError("distance must be a polars DataFrame.") 

59 

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

61 raise ValueError("Distance matrix index/columns must align with assets.") 

62 

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

64 raise ValueError("Number of leaves does not match number of assets.") 

65 

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

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

68 

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

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

71 """ 

72 from .plot import plot_dendrogram 

73 

74 return plot_dendrogram(self, **kwargs) 

75 

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

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

78 

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

80 the allocation core is imported lazily so importing the dendrogram module 

81 does not pull in the allocation module. 

82 

83 Yields: 

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

85 at that level. 

86 """ 

87 from .algos import one_over_n 

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 raise ValueError("ids must contain at least one node id.") 

115 if len(ids) == 1: 

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

117 

118 mid = len(ids) // 2 

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

120 left, next_id = _bisect_tree(left_ids, next_id) 

121 right, next_id = _bisect_tree(right_ids, next_id) 

122 next_id += 1 

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

124 

125 

126def _get_linkage(node: Cluster) -> list[list[float]]: 

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

128 links_list: list[list[float]] = [] 

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

130 if not isinstance(node.left, Cluster): 

131 raise TypeError("Expected left child to be a Cluster") # pragma: no cover 

132 if not isinstance(node.right, Cluster): 

133 raise TypeError("Expected right child to be a Cluster") # pragma: no cover 

134 links_list.extend(_get_linkage(node.left)) 

135 links_list.extend(_get_linkage(node.right)) 

136 links_list.append( 

137 [ 

138 float(node.left.value), 

139 float(node.right.value), 

140 float(node.size), 

141 float(len(node.left.leaves) + len(node.right.leaves)), 

142 ] 

143 ) 

144 return links_list 

145 

146 

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

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

149 

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

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

152 """ 

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

154 if bad: 

155 msg = ( 

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

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

158 ) 

159 raise ValueError(msg) 

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

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

162 raise ValueError(msg) 

163 

164 

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

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

167 

168 Raises: 

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

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

171 """ 

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

173 raise TypeError("Correlation matrix must be a polars DataFrame.") 

174 if len(cor.columns) < 2: 

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

176 raise ValueError(msg) 

177 _check_finite_correlations(cor, cor.to_numpy()) 

178 

179 

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

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

182 

183 Args: 

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

185 

186 Returns: 

187 Cluster: Equivalent node in our Cluster format. 

188 """ 

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

190 return Cluster(value=node.id, left=_to_cluster(node.left), right=_to_cluster(node.right)) 

191 return Cluster(value=node.id) 

192 

193 

194def build_tree( 

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

196) -> Dendrogram: 

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

198 

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

200 hierarchical clustering, and returns a Dendrogram object containing the 

201 resulting tree structure. 

202 

203 Args: 

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

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

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

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

208 - "average": average distance between all points 

209 - "ward": Ward variance minimization 

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

211 

212 Returns: 

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

214 - root: Root cluster node 

215 - linkage: Linkage matrix for plotting 

216 - assets: List of assets 

217 - method: Clustering method used 

218 - distance: Distance matrix 

219 

220 Examples: 

221 >>> import polars as pl 

222 >>> from pyhrp.dendrogram import build_tree 

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

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

225 >>> dg.root.leaf_count 

226 2 

227 """ 

228 _validate_correlation_matrix(cor) 

229 dist = _compute_distance_matrix(cor) 

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

231 

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

233 

234 # Apply bisection if requested 

235 if bisection: 

236 # Rebuild tree using bisection 

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

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

239 links = np.array(_get_linkage(root)) 

240 

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