API Reference¶

@dataclass(frozen=True)
class Dendrogram:
    """Container for hierarchical clustering dendrogram data and visualization.

    This class stores the results of hierarchical clustering and provides methods
    for accessing and visualizing the dendrogram structure.

    Attributes:
        root (Cluster): The root node of the hierarchical clustering tree
        assets (pd.Index): Index of assets included in the clustering
        linkage (np.ndarray | None): Linkage matrix in scipy format for plotting
        distance (np.ndarray | None): Distance matrix used for clustering
        method (str | None): Linkage method used for clustering
    """

    root: Cluster
    assets: pd.Index
    distance: pd.DataFrame | None = None
    linkage: np.ndarray | None = None
    method: str | None = None

    def __post_init__(self) -> None:
        """Validate dataclass fields after initialization.

        Ensures that the optional distance matrix, when provided, is a pandas
        DataFrame aligned with the asset order, and verifies that the number of
        leaves in the cluster tree matches the number of assets.
        """
        # ---- Optional: validate distance index/columns ----
        if self.distance is not None:
            if not isinstance(self.distance, pd.DataFrame):
                raise TypeError("distance must be a pandas DataFrame.")  # noqa: TRY003

            # Optionally check if distance matches assets
            if not self.distance.index.equals(pd.Index(self.assets)) or not self.distance.columns.equals(
                pd.Index(self.assets)
            ):
                raise ValueError("Distance matrix index/columns must align with assets.")  # noqa: TRY003

        # Check the number of leaves and assets
        if len(self.root.leaves) != len(self.assets):
            raise ValueError("Number of leaves does not match number of assets.")  # noqa: TRY003

    def plot(self, **kwargs: object) -> None:
        """Plot the dendrogram."""
        sch.dendrogram(self.linkage, leaf_rotation=90, leaf_font_size=8, labels=self.assets, **kwargs)

    @property
    def ids(self) -> list[int]:
        """Node values in the order left -> right as they appear in the dendrogram."""
        return [node.value for node in self.root.leaves]  # type: ignore[misc]

    @property
    def names(self) -> list[str]:
        """The asset names as induced by the order of ids."""
        return [self.assets[i] for i in self.ids]

class Cluster(Node):
    """Represents a cluster in the hierarchical clustering tree.

    Clusters are the nodes of the graphs we build.
    Each cluster is aware of the left and the right cluster
    it is connecting to. Each cluster also has an associated portfolio.

    Attributes:
        portfolio (Portfolio): The portfolio associated with this cluster
    """

    def __init__(self, value: int, left: Cluster | None = None, right: Cluster | None = None, **kwargs: Any) -> None:
        """Initialize a new Cluster.

        Args:
            value (int): The identifier for this cluster
            left (Cluster, optional): The left child cluster
            right (Cluster, optional): The right child cluster
            **kwargs: Additional arguments to pass to the parent class
        """
        super().__init__(value=value, left=left, right=right, **kwargs)
        self.portfolio = Portfolio()

    @property
    def is_leaf(self) -> bool:
        """Check if this cluster is a leaf node (has no children).

        Returns:
            bool: True if this is a leaf node, False otherwise
        """
        return self.left is None and self.right is None

    @property
    def leaves(self) -> list[Cluster]:
        """Get all reachable leaf nodes in the correct order.

        Note that the leaves method of the Node class implemented in BinaryTree
        is not respecting the 'correct' order of the nodes.

        Returns:
            list[Cluster]: List of all leaf nodes reachable from this cluster
        """
        if self.is_leaf:
            return [self]
        else:
            if self.left is None:
                raise ValueError("Expected left child to exist for non-leaf cluster")  # noqa: TRY003
            if self.right is None:
                raise ValueError("Expected right child to exist for non-leaf cluster")  # noqa: TRY003
            left_leaves: list[Cluster] = self.left.leaves  # type: ignore[assignment]
            right_leaves: list[Cluster] = self.right.leaves  # type: ignore[assignment]
            return left_leaves + right_leaves

@dataclass
class Portfolio:
    """Container for portfolio asset weights.

    This lightweight class stores and manipulates a mapping from asset names to
    their portfolio weights, and provides convenience helpers for analysis and
    visualization.

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

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

    @property
    def assets(self) -> list[str]:
        """List of asset names present in the portfolio.

        Returns:
            list[str]: Asset identifiers in insertion order (Python 3.7+ dict order).
        """
        return list(self._weights.keys())

    def variance(self, cov: pd.DataFrame) -> float:
        """Calculate the variance of the portfolio.

        Args:
            cov (pd.DataFrame): Covariance matrix

        Returns:
            float: Portfolio variance
        """
        c = cov[self.assets].loc[self.assets].values
        w = self.weights[self.assets].values
        return float(np.linalg.multi_dot((w, c, w)))

    def __getitem__(self, item: str) -> float:
        """Return the weight for a given asset.

        Args:
            item (str): Asset name/symbol.

        Returns:
            float: The weight associated with the asset.

        Raises:
            KeyError: If the asset is not present in the portfolio.
        """
        return self._weights[item]

    def __setitem__(self, key: str, value: float) -> None:
        """Set or update the weight for an asset.

        Args:
            key (str): Asset name/symbol.
            value (float): Portfolio weight for the asset.
        """
        self._weights[key] = value

    @property
    def weights(self) -> pd.Series:
        """Get all weights as a pandas Series.

        Returns:
            pd.Series: Series of weights indexed by assets
        """
        return pd.Series(self._weights, name="Weights").sort_index()

    def plot(self, names: list[str]) -> Axes:
        """Plot the portfolio weights.

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

        Returns:
            matplotlib.axes.Axes: The plot axes
        """
        a = self.weights.loc[names]

        ax = a.plot(kind="bar", color="skyblue")

        # Set x-axis labels and rotations
        ax.set_xticklabels(names, rotation=90, fontsize=8)
        return ax

class Node:
    """A binary tree node with left and right children.

    This class implements the minimal functionality needed from the binarytree.Node class
    that is used in the pyhrp package.

    Attributes:
        value: The value of the node
        left: The left child node
        right: The right child node
    """

    def __init__(self, value: NodeValue, left: Node | None = None, right: Node | None = None):
        """Initialize a new Node.

        Args:
            value: The value of the node
            left: The left child node
            right: The right child node
        """
        self.value = value
        self.left = left
        self.right = right

    @property
    def is_leaf(self) -> bool:
        """Check if this node is a leaf node (has no children).

        Returns:
            bool: True if this is a leaf node, False otherwise
        """
        return self.left is None and self.right is None

    @property
    def leaves(self) -> list[Node]:
        """Get all leaf nodes in the tree rooted at this node.

        Returns:
            List[Node]: List of all leaf nodes
        """
        if self.is_leaf:
            return [self]

        result = []
        if self.left:
            result.extend(self.left.leaves)
        if self.right:
            result.extend(self.right.leaves)

        return result

    @property
    def levels(self) -> list[list[Node]]:
        """Get nodes by level in the tree.

        Returns:
            List[List[Node]]: List of lists of nodes at each level
        """
        result = []
        current_level = [self]

        while current_level:
            result.append(current_level)
            next_level = []

            for node in current_level:
                if node.left:
                    next_level.append(node.left)
                if node.right:
                    next_level.append(node.right)

            current_level = next_level

        return result

    @property
    def leaf_count(self) -> int:
        """Count the number of leaf nodes in the tree.

        Returns:
            int: Number of leaf nodes
        """
        return len(self.leaves)

    @property
    def size(self) -> int:
        """Count the total number of nodes in the tree.

        Returns:
            int: Total number of nodes
        """
        size = 1  # Count this node
        if self.left:
            size += self.left.size
        if self.right:
            size += self.right.size
        return size

    def __iter__(self) -> Iterator[Node]:
        """Iterate through all nodes in the tree in level-order.

        Returns:
            Iterator[Node]: Iterator over all nodes
        """
        queue: deque[Node] = deque([self])
        while queue:
            node = queue.popleft()
            yield node
            if node.left:
                queue.append(node.left)
            if node.right:
                queue.append(node.right)