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

1"""Data structures for hierarchical risk parity portfolio optimization. 

2 

3This module defines the core data structures used in the hierarchical risk parity algorithm: 

4- Portfolio: Manages a collection of asset weights (strings identify assets) 

5- Cluster: Represents a node in the hierarchical clustering tree 

6""" 

7 

8from __future__ import annotations 

9 

10from dataclasses import dataclass, field 

11from typing import TYPE_CHECKING 

12 

13import numpy as np 

14import polars as pl 

15 

16from .treelib import Node 

17 

18if TYPE_CHECKING: 

19 import plotly.graph_objects as go 

20 

21__all__ = ["Cluster", "Portfolio"] 

22 

23 

24@dataclass 

25class Portfolio: 

26 """Container for portfolio asset weights. 

27 

28 This lightweight class stores and manipulates a mapping from asset names to 

29 their portfolio weights, and provides convenience helpers for analysis and 

30 visualization. 

31 

32 Attributes: 

33 _weights (dict[str, float]): Internal mapping from asset symbol to weight. 

34 """ 

35 

36 _weights: dict[str, float] = field(default_factory=dict) 

37 

38 @property 

39 def assets(self) -> list[str]: 

40 """List of asset names present in the portfolio. 

41 

42 Returns: 

43 list[str]: Asset identifiers in insertion order (Python 3.7+ dict order). 

44 """ 

45 return list(self._weights.keys()) 

46 

47 def variance(self, cov: pl.DataFrame) -> float: 

48 """Calculate the variance of the portfolio. 

49 

50 Args: 

51 cov (pl.DataFrame): Covariance matrix where columns and rows correspond 

52 to assets in the same order as columns list. 

53 

54 Returns: 

55 float: Portfolio variance 

56 """ 

57 assets = self.assets 

58 index = {name: i for i, name in enumerate(cov.columns)} 

59 row_indices = [index[a] for a in assets] 

60 cov_matrix = cov.to_numpy() 

61 c = cov_matrix[np.ix_(row_indices, row_indices)] 

62 w = np.array([self._weights[a] for a in assets]) 

63 return float(w @ c @ w) 

64 

65 def __getitem__(self, item: str) -> float: 

66 """Return the weight for a given asset. 

67 

68 Args: 

69 item (str): Asset name/symbol. 

70 

71 Returns: 

72 float: The weight associated with the asset. 

73 

74 Raises: 

75 KeyError: If the asset is not present in the portfolio. 

76 """ 

77 return self._weights[item] 

78 

79 def __setitem__(self, key: str, value: float) -> None: 

80 """Set or update the weight for an asset. 

81 

82 Args: 

83 key (str): Asset name/symbol. 

84 value (float): Portfolio weight for the asset. 

85 """ 

86 self._weights[key] = value 

87 

88 @property 

89 def weights(self) -> dict[str, float]: 

90 """Get all weights as a dict sorted alphabetically by asset name. 

91 

92 Returns: 

93 dict[str, float]: Mapping from asset name to weight, sorted by name. 

94 """ 

95 return dict(sorted(self._weights.items())) 

96 

97 def plot(self, names: list[str]) -> go.Figure: 

98 """Plot the portfolio weights as a bar chart. 

99 

100 Args: 

101 names (list[str]): List of asset names to include in the plot 

102 

103 Returns: 

104 go.Figure: The plotly figure 

105 

106 Note: 

107 The plotly dependency is imported lazily so importing the allocation 

108 core (which pulls in this module) stays plotly-free. 

109 """ 

110 import plotly.graph_objects as go 

111 

112 w = self.weights 

113 values = [w[n] for n in names] 

114 fig = go.Figure(go.Bar(x=names, y=values, marker_color="steelblue")) 

115 fig.update_layout(xaxis={"tickangle": -90}) 

116 return fig 

117 

118 

119class Cluster(Node[int]): 

120 """Represents a cluster in the hierarchical clustering tree. 

121 

122 Clusters are the nodes of the graphs we build. 

123 Each cluster is aware of the left and the right cluster 

124 it is connecting to. Each cluster also has an associated portfolio. 

125 

126 Attributes: 

127 portfolio (Portfolio): The portfolio associated with this cluster 

128 """ 

129 

130 def __init__(self, value: int, left: Cluster | None = None, right: Cluster | None = None) -> None: 

131 """Initialize a new Cluster. 

132 

133 Args: 

134 value (int): The identifier for this cluster 

135 left (Cluster, optional): The left child cluster 

136 right (Cluster, optional): The right child cluster 

137 """ 

138 super().__init__(value=value, left=left, right=right) 

139 self.portfolio = Portfolio() 

140 

141 # Override narrows the return type to list[Cluster] and validates tree integrity; 

142 # the traversal order (left to right) matches Node.leaves. 

143 @property 

144 def leaves(self) -> list[Cluster]: 

145 """Get all reachable leaf nodes in left-to-right dendrogram order. 

146 

147 Returns: 

148 list[Cluster]: List of all leaf nodes reachable from this cluster 

149 """ 

150 if self.is_leaf: 

151 return [self] 

152 left, right = self._child_clusters() 

153 return left.leaves + right.leaves 

154 

155 def _child_clusters(self) -> tuple[Cluster, Cluster]: 

156 """Return the validated (left, right) child clusters of a non-leaf node. 

157 

158 Raises: 

159 ValueError: If either child is missing on a non-leaf cluster. 

160 TypeError: If either child is not a Cluster. 

161 """ 

162 if self.left is None: 

163 raise ValueError("Expected left child to exist for non-leaf cluster") 

164 if self.right is None: 

165 raise ValueError("Expected right child to exist for non-leaf cluster") 

166 if not isinstance(self.left, Cluster): 

167 raise TypeError(f"Expected left child to be a Cluster for node {self.value}") 

168 if not isinstance(self.right, Cluster): 

169 raise TypeError(f"Expected right child to be a Cluster for node {self.value}") 

170 return self.left, self.right