Coverage for src/pyhrp/cluster.py: 100%
66 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-09-15 05:20 +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 dataclasses import dataclass, field
11from typing import TYPE_CHECKING
13import numpy as np
14import polars as pl
16from .treelib import Node
18if TYPE_CHECKING:
19 import plotly.graph_objects as go
21__all__ = ["Cluster", "Portfolio"]
24@dataclass
25class Portfolio:
26 """Container for portfolio asset weights.
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.
32 Attributes:
33 _weights (dict[str, float]): Internal mapping from asset symbol to weight.
34 """
36 _weights: dict[str, float] = field(default_factory=dict)
38 @property
39 def assets(self) -> list[str]:
40 """List of asset names present in the portfolio.
42 Returns:
43 list[str]: Asset identifiers in insertion order (Python 3.7+ dict order).
44 """
45 return list(self._weights.keys())
47 def variance(self, cov: pl.DataFrame) -> float:
48 """Calculate the variance of the portfolio.
50 Args:
51 cov (pl.DataFrame): Covariance matrix where columns and rows correspond
52 to assets in the same order as columns list.
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)
65 def __getitem__(self, item: str) -> float:
66 """Return the weight for a given asset.
68 Args:
69 item (str): Asset name/symbol.
71 Returns:
72 float: The weight associated with the asset.
74 Raises:
75 KeyError: If the asset is not present in the portfolio.
76 """
77 return self._weights[item]
79 def __setitem__(self, key: str, value: float) -> None:
80 """Set or update the weight for an asset.
82 Args:
83 key (str): Asset name/symbol.
84 value (float): Portfolio weight for the asset.
85 """
86 self._weights[key] = value
88 @property
89 def weights(self) -> dict[str, float]:
90 """Get all weights as a dict sorted alphabetically by asset name.
92 Returns:
93 dict[str, float]: Mapping from asset name to weight, sorted by name.
94 """
95 return dict(sorted(self._weights.items()))
97 def plot(self, names: list[str]) -> go.Figure:
98 """Plot the portfolio weights as a bar chart.
100 Args:
101 names (list[str]): List of asset names to include in the plot
103 Returns:
104 go.Figure: The plotly figure
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
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
119class Cluster(Node[int]):
120 """Represents a cluster in the hierarchical clustering tree.
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.
126 Attributes:
127 portfolio (Portfolio): The portfolio associated with this cluster
128 """
130 def __init__(self, value: int, left: Cluster | None = None, right: Cluster | None = None) -> None:
131 """Initialize a new Cluster.
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()
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.
147 Returns:
148 list[Cluster]: List of all leaf nodes reachable from this cluster
149 """
150 # Iterative for the same reason as Node.leaves, and still validating each
151 # non-leaf node through _child_clusters() so a malformed tree raises here
152 # rather than yielding a silently short leaf list.
153 result: list[Cluster] = []
154 stack: list[Cluster] = [self]
155 while stack:
156 node = stack.pop()
157 if node.is_leaf:
158 result.append(node)
159 continue
160 left, right = node._child_clusters()
161 # Right first, so the left subtree is emitted first.
162 stack.append(right)
163 stack.append(left)
164 return result
166 def _child_clusters(self) -> tuple[Cluster, Cluster]:
167 """Return the validated (left, right) child clusters of a non-leaf node.
169 Raises:
170 ValueError: If either child is missing on a non-leaf cluster.
171 TypeError: If either child is not a Cluster.
172 """
173 if self.left is None:
174 msg = "Expected left child to exist for non-leaf cluster"
175 raise ValueError(msg)
176 if self.right is None:
177 msg = "Expected right child to exist for non-leaf cluster"
178 raise ValueError(msg)
179 if not isinstance(self.left, Cluster):
180 msg = f"Expected left child to be a Cluster for node {self.value}"
181 raise TypeError(msg)
182 if not isinstance(self.right, Cluster):
183 msg = f"Expected right child to be a Cluster for node {self.value}"
184 raise TypeError(msg)
185 return self.left, self.right