API Reference¶

def build_tree(
    cor: pl.DataFrame, method: Literal["single", "complete", "average", "ward"] = "ward", bisection: bool = False
) -> Dendrogram:
    """Build hierarchical cluster tree from correlation matrix.

    This function converts a correlation matrix to a distance matrix, performs
    hierarchical clustering, and returns a Dendrogram object containing the
    resulting tree structure.

    Args:
        cor (pl.DataFrame): Correlation matrix of asset returns (columns are assets)
        method (Literal["single", "complete", "average", "ward"]): Linkage method for hierarchical clustering
            - "single": minimum distance between points (nearest neighbor)
            - "complete": maximum distance between points (furthest neighbor)
            - "average": average distance between all points
            - "ward": Ward variance minimization
        bisection (bool): Whether to use bisection method for tree construction

    Returns:
        Dendrogram: Object containing the hierarchical clustering tree, with:
            - root: Root cluster node
            - linkage: Linkage matrix for plotting
            - assets: List of assets
            - method: Clustering method used
            - distance: Distance matrix

    Examples:
        >>> import polars as pl
        >>> from pyhrp.dendrogram import build_tree
        >>> cor = pl.DataFrame({"A": [1.0, 0.5], "B": [0.5, 1.0]})
        >>> dg = build_tree(cor, method="ward")
        >>> dg.root.leaf_count
        2
    """
    _validate_correlation_matrix(cor)
    dist = _compute_distance_matrix(cor)
    links = sch.linkage(ssd.squareform(dist.to_numpy(), checks=False), method=method)

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

    # Apply bisection if requested
    if bisection:
        # Rebuild tree using bisection
        leaf_ids: list[int] = [int(node.value) for node in root.leaves]
        root, _ = _bisect_tree(ids=leaf_ids, next_id=max(leaf_ids))
        links = np.array(_get_linkage(root))

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

def compute_corr(df: pl.DataFrame) -> pl.DataFrame:
    """Compute correlation matrix from a DataFrame of returns."""
    cols = df.columns
    corr = np.atleast_2d(np.corrcoef(df.to_numpy().T))
    return pl.DataFrame(dict(zip(cols, corr, strict=True)))

def compute_cov(df: pl.DataFrame) -> pl.DataFrame:
    """Compute covariance matrix from a DataFrame of returns."""
    cols = df.columns
    cov = np.atleast_2d(np.cov(df.to_numpy().T))
    return pl.DataFrame(dict(zip(cols, cov, strict=True)))

def hrp(
    prices: pl.DataFrame,
    node: Cluster | None = None,
    method: Literal["single", "complete", "average", "ward"] = "ward",
    bisection: bool = False,
) -> Cluster:
    """Compute the hierarchical risk parity portfolio weights.

    This is the main entry point for the HRP algorithm. It calculates returns from prices,
    builds a hierarchical clustering tree if not provided, and applies risk parity weights.

    Args:
        prices (pl.DataFrame): Asset price time series (columns are assets, rows are dates)
        node (Cluster, optional): Root node of the hierarchical clustering tree.
            If None, a tree will be built from the correlation matrix.
        method (Literal["single", "complete", "average", "ward"]): Linkage method to use for distance calculation
            - "single": minimum distance between points (nearest neighbor)
            - "complete": maximum distance between points (furthest neighbor)
            - "average": average distance between all points
            - "ward": Ward variance minimization
        bisection (bool): Whether to use bisection method for tree construction

    Returns:
        Cluster: The root cluster with portfolio weights assigned according to HRP

    Examples:
        >>> import polars as pl
        >>> from pyhrp.hrp import hrp
        >>> prices = pl.DataFrame({"A": [100.0, 101.0, 99.0, 102.0], "B": [50.0, 51.0, 49.0, 52.0]})
        >>> root = hrp(prices, method="ward")
        >>> round(sum(root.portfolio.weights.values()), 6)
        1.0
    """
    returns = _returns(prices)
    cov = compute_cov(returns)
    cor = compute_corr(returns)
    node = node or build_tree(cor, method=method, bisection=bisection).root

    return risk_parity(root=node, cov=cov)

def schur_hrp(
    prices: pl.DataFrame,
    node: Cluster | None = None,
    method: Literal["single", "complete", "average", "ward"] = "ward",
    bisection: bool = False,
    gamma: float = 0.5,
) -> Cluster:
    """Compute Schur Complementary Allocation portfolio weights.

    Extends HRP by augmenting each sub-covariance block with off-diagonal information
    via Schur complements before splitting risk between clusters. Introduced by Peter Cotton
    (arXiv:2411.05807). At gamma=0 this is identical to HRP; at gamma=1 it recovers the
    global minimum-variance portfolio through the same recursive hierarchy.

    Args:
        prices (pl.DataFrame): Asset price time series (columns are assets, rows are dates)
        node (Cluster, optional): Root node of the hierarchical clustering tree.
            If None, a tree will be built from the correlation matrix.
        method (Literal["single", "complete", "average", "ward"]): Linkage method for clustering
        bisection (bool): Whether to use bisection method for tree construction
        gamma (float): Schur interpolation parameter in [0, 1].
            0 recovers standard HRP; 1 recovers minimum-variance portfolio.

    Returns:
        Cluster: The root cluster with portfolio weights assigned

    Examples:
        >>> import polars as pl
        >>> from pyhrp.hrp import schur_hrp
        >>> prices = pl.DataFrame({"A": [100.0, 101.0, 99.0, 102.0], "B": [50.0, 51.0, 49.0, 52.0]})
        >>> root = schur_hrp(prices, method="ward", gamma=0.5)
        >>> round(sum(root.portfolio.weights.values()), 6)
        1.0
    """
    returns = _returns(prices)
    cov = compute_cov(returns)
    cor = compute_corr(returns)
    node = node or build_tree(cor, method=method, bisection=bisection).root

    return schur_risk_parity(root=node, cov=cov, gamma=gamma)

def one_over_n(root: Cluster, assets: list[str]) -> Generator[tuple[int, Portfolio]]:
    """Generate 1/N (equal-weight) portfolios one tree level at a time.

    This implements a hierarchical 1/N strategy where weights are distributed
    equally among the leaves of each cluster, and the weight budget halves at
    each successive level of the tree.

    Unlike :func:`risk_parity` and :func:`schur_risk_parity` — which rebuild a
    single final allocation and return the root ``Cluster`` — this allocator is
    intentionally a **generator**: its purpose is to expose the equal-weight
    allocation as the tree deepens, yielding one portfolio per level. It shares
    the sibling input contract (a ``Cluster`` tree plus the asset names) and, like
    them, does not mutate the tree: weights accumulate in a local buffer, so a
    leaf that terminates at a shallow level keeps its weight in the deeper levels
    (each yielded portfolio is therefore a complete allocation over all assets),
    and re-running on the same tree yields an identical sequence.

    Args:
        root (Cluster): The root node of the cluster tree.
        assets (list[str]): Asset names; a leaf's value indexes into this list.

    Yields:
        tuple[int, Portfolio]: The level number and the (cumulative) equal-weight
        portfolio at that level.

    Examples:
        >>> import polars as pl
        >>> from pyhrp.hrp import build_tree
        >>> from pyhrp.algos import one_over_n
        >>> cor = pl.DataFrame({"A": [1.0, 0.3], "B": [0.3, 1.0]})
        >>> dg = build_tree(cor, method="ward")
        >>> levels = list(one_over_n(dg.root, dg.assets))
        >>> len(levels) > 0
        True
    """
    # Accumulate into a local buffer so the input tree is never mutated.
    portfolio = Portfolio()

    # Initial weight to distribute
    w: float = 1.0

    # Process each level of the tree
    for n, level in enumerate(root.levels):
        for node in level:
            # Distribute weight equally among all leaves in this node
            for leaf in node.leaves:
                portfolio[assets[leaf.value]] = w / node.leaf_count

        # Reduce weight for the next level
        w *= 0.5

        # Yield the current level number and a deep copy of the portfolio
        yield n, deepcopy(portfolio)

def risk_parity(root: Cluster, cov: pl.DataFrame) -> Cluster:
    """Compute hierarchical risk parity weights for a cluster tree.

    This is the main algorithm for hierarchical risk parity. It recursively
    traverses the cluster tree and assigns weights to each node based on
    the risk parity principle.

    Note:
        The tree is modified in place: the portfolio of every node is rebuilt
        from scratch, so the function is idempotent and a tree can be reused
        with a different covariance matrix.

    Args:
        root (Cluster): The root node of the cluster tree
        cov (pl.DataFrame): Covariance matrix of asset returns

    Returns:
        Cluster: The root node with portfolio weights assigned

    Examples:
        >>> import polars as pl
        >>> from pyhrp.cluster import Cluster
        >>> from pyhrp.algos import risk_parity
        >>> cov = pl.DataFrame({"A": [4.0, 0.0], "B": [0.0, 1.0]})
        >>> root = Cluster(2, left=Cluster(0), right=Cluster(1))
        >>> cluster = risk_parity(root=root, cov=cov)
        >>> round(cluster.portfolio["B"], 1)
        0.8
    """

    def node_variances(left: Cluster, right: Cluster, cov_np: np.ndarray, index: dict[str, int]) -> tuple[float, float]:
        """Plain block variance of each child sub-portfolio."""
        return (
            _block_variance(left.portfolio, cov_np, index),
            _block_variance(right.portfolio, cov_np, index),
        )

    return _allocate_with(root, cov, node_variances)

def schur_risk_parity(root: Cluster, cov: pl.DataFrame, gamma: float = 0.5) -> Cluster:
    """Compute Schur Complementary Allocation weights for a cluster tree.

    An extension of HRP introduced by Peter Cotton (arXiv:2411.05807) that augments
    sub-covariance matrices with off-diagonal block information via Schur complements.
    At gamma=0 this recovers standard HRP; at gamma=1 it recovers the minimum-variance
    portfolio through the same recursive structure.

    Note:
        The tree is modified in place: the portfolio of every node is rebuilt
        from scratch, so the function is idempotent and a tree can be reused
        with a different covariance matrix or gamma.

    Args:
        root (Cluster): The root node of the cluster tree
        cov (pl.DataFrame): Covariance matrix of asset returns
        gamma (float): Interpolation parameter in [0, 1]. 0 = HRP, 1 = minimum variance.

    Returns:
        Cluster: The root node with portfolio weights assigned

    Raises:
        ValueError: If gamma is outside the interval [0, 1].

    Examples:
        >>> import polars as pl
        >>> from pyhrp.cluster import Cluster
        >>> from pyhrp.algos import schur_risk_parity
        >>> cov = pl.DataFrame({"A": [4.0, 0.0], "B": [0.0, 1.0]})
        >>> root = Cluster(2, left=Cluster(0), right=Cluster(1))
        >>> cluster = schur_risk_parity(root=root, cov=cov, gamma=0.5)
        >>> round(cluster.portfolio["B"], 1)
        0.8
    """
    if not 0.0 <= gamma <= 1.0:
        msg = f"gamma must be in [0, 1], got {gamma}"
        raise ValueError(msg)

    def node_variances(left: Cluster, right: Cluster, cov_np: np.ndarray, index: dict[str, int]) -> tuple[float, float]:
        """Schur-augmented block variance of each child, conditioned on the other."""
        li = [index[a] for a in left.portfolio.assets]
        ri = [index[a] for a in right.portfolio.assets]

        a_mat = cov_np[np.ix_(li, li)]
        b_mat = cov_np[np.ix_(li, ri)]
        d_mat = cov_np[np.ix_(ri, ri)]

        w_left = np.array([left.portfolio[a] for a in left.portfolio.assets])
        w_right = np.array([right.portfolio[a] for a in right.portfolio.assets])

        # Schur-augmented blocks: condition each group on the other
        a_aug = a_mat - gamma * (b_mat @ _solve(d_mat, b_mat.T))
        d_aug = d_mat - gamma * (b_mat.T @ _solve(a_mat, b_mat))

        v_left = float(w_left @ a_aug @ w_left)
        v_right = float(w_right @ d_aug @ w_right)
        return v_left, v_right

    return _allocate_with(root, cov, node_variances)

def plot_dendrogram(dendrogram: Dendrogram, **kwargs: object) -> go.Figure:
    """Build and return a plotly dendrogram figure for a Dendrogram.

    Args:
        dendrogram (Dendrogram): Clustering result carrying the linkage matrix
            and asset labels to render.
        **kwargs (object): Extra keyword arguments forwarded to
            ``scipy.cluster.hierarchy.dendrogram`` (e.g. ``color_threshold``).

    Returns:
        go.Figure: A plotly figure drawing the dendrogram as line traces.

    Raises:
        ValueError: If the dendrogram has no linkage matrix to plot.

    Examples:
        >>> import polars as pl
        >>> from pyhrp.hrp import build_tree
        >>> from pyhrp.plot import plot_dendrogram
        >>> cor = pl.DataFrame({"A": [1.0, 0.5], "B": [0.5, 1.0]})
        >>> fig = plot_dendrogram(build_tree(cor, method="ward"))
        >>> type(fig).__name__
        'Figure'
    """
    if dendrogram.linkage is None:
        msg = "Dendrogram has no linkage matrix to plot."
        raise ValueError(msg)
    ddata = sch.dendrogram(dendrogram.linkage, labels=dendrogram.assets, no_plot=True, **kwargs)
    fig = go.Figure()
    for xs, ys in zip(ddata["icoord"], ddata["dcoord"], strict=False):
        fig.add_trace(go.Scatter(x=xs, y=ys, mode="lines", line={"color": "steelblue"}, showlegend=False))
    n = len(dendrogram.assets)
    fig.update_layout(
        xaxis={
            "tickmode": "array",
            "tickvals": [5 + 10 * i for i in range(n)],
            "ticktext": ddata["ivl"],
            "tickangle": -90,
        },
    )
    return fig