Coverage for src/pyhrp/cluster.py: 100%
124 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-11-03 05:36 +0000
1"""Data structures for hierarchical risk parity portfolio optimization.
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"""
8from __future__ import annotations
10from collections.abc import Sequence
11from dataclasses import dataclass, field
12from typing import TYPE_CHECKING
14import numpy as np
15import polars as pl
17from .treelib import Node
19if TYPE_CHECKING:
20 import plotly.graph_objects as go
22__all__ = ["Cluster", "Portfolio"]
24# One non-leaf node of a laid-out tree: the node and the half-open slices
25# [lo, mid) and [mid, hi) of the leaf order its left and right subtrees cover.
26Span = tuple["Cluster", int, int, int]
29@dataclass
30class Portfolio:
31 """Container for portfolio asset weights.
33 This lightweight class stores and manipulates a mapping from asset names to
34 their portfolio weights, and provides convenience helpers for analysis and
35 visualization.
37 Attributes:
38 _weights (dict[str, float]): Internal mapping from asset symbol to weight.
39 """
41 _weights: dict[str, float] = field(default_factory=dict)
43 @property
44 def assets(self) -> list[str]:
45 """List of asset names present in the portfolio.
47 Returns:
48 list[str]: Asset identifiers in insertion order (Python 3.7+ dict order).
49 """
50 return list(self._weights.keys())
52 def variance(self, cov: pl.DataFrame) -> float:
53 """Calculate the variance of the portfolio.
55 Args:
56 cov (pl.DataFrame): Covariance matrix where columns and rows correspond
57 to assets in the same order as columns list.
59 Returns:
60 float: Portfolio variance
61 """
62 assets = self.assets
63 index = {name: i for i, name in enumerate(cov.columns)}
64 row_indices = [index[a] for a in assets]
65 cov_matrix = cov.to_numpy()
66 c = cov_matrix[np.ix_(row_indices, row_indices)]
67 w = np.array([self._weights[a] for a in assets])
68 return float(w @ c @ w)
70 def __getitem__(self, item: str) -> float:
71 """Return the weight for a given asset.
73 Args:
74 item (str): Asset name/symbol.
76 Returns:
77 float: The weight associated with the asset.
79 Raises:
80 KeyError: If the asset is not present in the portfolio.
81 """
82 return self._weights[item]
84 def __setitem__(self, key: str, value: float) -> None:
85 """Set or update the weight for an asset.
87 Args:
88 key (str): Asset name/symbol.
89 value (float): Portfolio weight for the asset.
90 """
91 self._weights[key] = value
93 @property
94 def weights(self) -> dict[str, float]:
95 """Get all weights as a dict sorted alphabetically by asset name.
97 Returns:
98 dict[str, float]: Mapping from asset name to weight, sorted by name.
99 """
100 return dict(sorted(self._weights.items()))
102 def plot(self, names: list[str]) -> go.Figure:
103 """Plot the portfolio weights as a bar chart.
105 Args:
106 names (list[str]): List of asset names to include in the plot
108 Returns:
109 go.Figure: The plotly figure
111 Note:
112 The plotly dependency is imported lazily so importing the allocation
113 core (which pulls in this module) stays plotly-free.
114 """
115 import plotly.graph_objects as go
117 w = self.weights
118 values = [w[n] for n in names]
119 fig = go.Figure(go.Bar(x=names, y=values, marker_color="steelblue"))
120 fig.update_layout(xaxis={"tickangle": -90})
121 return fig
124class Cluster(Node[int]):
125 """Represents a cluster in the hierarchical clustering tree.
127 Clusters are the nodes of the graphs we build.
128 Each cluster is aware of the left and the right cluster
129 it is connecting to. Each cluster also has an associated portfolio.
131 Attributes:
132 portfolio (Portfolio): The portfolio associated with this cluster
134 After an allocation a node does not hold its portfolio, only its split: the
135 share of its weight that goes to the left child. The portfolio is built from
136 the splits below the node when it is first read, and then kept. Holding every
137 node's portfolio would cost the sum of all cluster sizes, which is quadratic in
138 the number of assets for the chain-like trees single linkage builds.
139 """
141 def __init__(self, value: int, left: Cluster | None = None, right: Cluster | None = None) -> None:
142 """Initialize a new Cluster.
144 Args:
145 value (int): The identifier for this cluster
146 left (Cluster, optional): The left child cluster
147 right (Cluster, optional): The right child cluster
148 """
149 super().__init__(value=value, left=left, right=right)
150 self._portfolio: Portfolio | None = Portfolio()
151 self._share = 1.0
152 self._assets: Sequence[str] = ()
154 @property
155 def portfolio(self) -> Portfolio:
156 """The portfolio of this cluster, built from the allocation's splits on first access."""
157 if self._portfolio is None:
158 leaves, spans = self._layout()
159 w = np.ones(len(leaves))
160 _apply_shares(w, spans)
161 self._portfolio = _portfolio(leaves, w, self._assets)
162 return self._portfolio
164 @portfolio.setter
165 def portfolio(self, value: Portfolio) -> None:
166 """Set the portfolio explicitly."""
167 self._portfolio = value
169 def _set_allocation(self, share: float, assets: Sequence[str]) -> None:
170 """Record this node's left share and drop its portfolio, to be rebuilt on access."""
171 self._share = share
172 self._assets = assets
173 self._portfolio = None
175 def _layout(self) -> tuple[list[Cluster], list[Span]]:
176 """Lay the subtree out over its leaf order.
178 Returns the leaves left to right and, for every non-leaf node in
179 post-order (children before parents), the slices of that leaf order its
180 two subtrees cover. Iterative and linear in the subtree size; every
181 non-leaf node is validated before anything is returned.
182 """
183 preorder: list[Cluster] = []
184 stack: list[Cluster] = [self]
185 while stack:
186 node = stack.pop()
187 preorder.append(node)
188 if not node.is_leaf:
189 left, right = node._child_clusters()
190 stack.append(right)
191 stack.append(left)
193 # Leaf counts bottom-up, then each node's first leaf position top-down.
194 count: dict[int, int] = {}
195 for node in reversed(preorder):
196 if node.is_leaf:
197 count[id(node)] = 1
198 else:
199 left, right = node._child_clusters()
200 count[id(node)] = count[id(left)] + count[id(right)]
201 start = {id(self): 0}
202 leaves: list[Cluster] = []
203 spans: list[Span] = []
204 for node in preorder:
205 lo = start[id(node)]
206 if node.is_leaf:
207 leaves.append(node)
208 else:
209 left, right = node._child_clusters()
210 mid = lo + count[id(left)]
211 start[id(left)] = lo
212 start[id(right)] = mid
213 spans.append((node, lo, mid, lo + count[id(node)]))
214 spans.reverse()
215 return leaves, spans
217 # Override narrows the return type to list[Cluster] and validates tree integrity;
218 # the traversal order (left to right) matches Node.leaves.
219 @property
220 def leaves(self) -> list[Cluster]:
221 """Get all reachable leaf nodes in left-to-right dendrogram order.
223 Returns:
224 list[Cluster]: List of all leaf nodes reachable from this cluster
225 """
226 # Iterative for the same reason as Node.leaves, and still validating each
227 # non-leaf node through _child_clusters() so a malformed tree raises here
228 # rather than yielding a silently short leaf list.
229 result: list[Cluster] = []
230 stack: list[Cluster] = [self]
231 while stack:
232 node = stack.pop()
233 if node.is_leaf:
234 result.append(node)
235 continue
236 left, right = node._child_clusters()
237 # Right first, so the left subtree is emitted first.
238 stack.append(right)
239 stack.append(left)
240 return result
242 def _child_clusters(self) -> tuple[Cluster, Cluster]:
243 """Return the validated (left, right) child clusters of a non-leaf node.
245 Raises:
246 ValueError: If either child is missing on a non-leaf cluster.
247 TypeError: If either child is not a Cluster.
248 """
249 if self.left is None:
250 msg = "Expected left child to exist for non-leaf cluster"
251 raise ValueError(msg)
252 if self.right is None:
253 msg = "Expected right child to exist for non-leaf cluster"
254 raise ValueError(msg)
255 if not isinstance(self.left, Cluster):
256 msg = f"Expected left child to be a Cluster for node {self.value}"
257 raise TypeError(msg)
258 if not isinstance(self.right, Cluster):
259 msg = f"Expected right child to be a Cluster for node {self.value}"
260 raise TypeError(msg)
261 return self.left, self.right
264def _apply_shares(w: np.ndarray, spans: list[Span]) -> None:
265 """Scale ``w`` bottom-up by the left share of every node, in place.
267 ``w`` starts at one per leaf; afterwards it holds the subtree's portfolio in
268 leaf order. The products are formed bottom-up, the order in which the
269 allocators form them.
270 """
271 for node, lo, mid, hi in spans:
272 w[lo:mid] *= node._share
273 w[mid:hi] *= 1.0 - node._share
276def _portfolio(leaves: list[Cluster], w: np.ndarray, assets: Sequence[str]) -> Portfolio:
277 """The Portfolio holding weight ``w[i]`` for the asset of ``leaves[i]``, in leaf order."""
278 portfolio = Portfolio()
279 for leaf, weight in zip(leaves, w.tolist(), strict=True):
280 portfolio[assets[int(leaf.value)]] = weight
281 return portfolio