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.hrp 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
    """
    if not isinstance(cor, pl.DataFrame):
        raise TypeError("Correlation matrix must be a polars DataFrame.")  # noqa: TRY003
    if len(cor.columns) < 2:
        msg = "Correlation matrix must contain at least two assets."
        raise ValueError(msg)
    c = cor.to_numpy()
    bad = [col for col, diag in zip(cor.columns, np.diagonal(c), strict=True) if not np.isfinite(diag)]
    if bad:
        msg = (
            f"Correlation matrix contains non-finite values for assets {bad}; "
            "constant (zero-variance) price series produce NaN correlations."
        )
        raise ValueError(msg)
    if not np.isfinite(c).all():
        msg = "Correlation matrix contains non-finite values."
        raise ValueError(msg)
    dist = _compute_distance_matrix(cor)
    links = sch.linkage(ssd.squareform(dist.to_numpy(), checks=False), method=method)

    # Convert scipy tree to our Cluster format
    def to_cluster(node: sch.ClusterNode) -> Cluster:
        """Convert a scipy ClusterNode to our Cluster format.

        Args:
            node (sch.ClusterNode): A node from scipy's hierarchical clustering

        Returns:
            Cluster: Equivalent node in our Cluster format
        """
        if node.left is not None and node.right is not None:
            left = to_cluster(node.left)
            right = to_cluster(node.right)
            return Cluster(value=node.id, left=left, right=right)
        return Cluster(value=node.id)

    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(dendrogram: Dendrogram) -> Generator[tuple[int, Portfolio]]:
    """Generate portfolios using the 1/N (equal weight) strategy at each tree level.

    This function implements a hierarchical 1/N strategy where weights are
    distributed equally among assets within each cluster at each level of the tree.
    The weight assigned to each cluster decreases by half at each level.

    Args:
        dendrogram: A dendrogram object containing the hierarchical clustering tree
                   and the list of assets

    Yields:
        tuple[int, Portfolio]: A tuple containing the level number and the 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))
        >>> len(levels) > 0
        True
    """
    root = dendrogram.root
    assets = dendrogram.assets

    # 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:
                root.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(root.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
    """
    cov_np = cov.to_numpy()
    index = {name: i for i, name in enumerate(cov.columns)}

    def combine(cluster: Cluster) -> Cluster:
        """Combine the child portfolios of a cluster via inverse-variance split."""
        left, right = _children(cluster)
        v_left = _block_variance(left.portfolio, cov_np, index)
        v_right = _block_variance(right.portfolio, cov_np, index)
        return _split(cluster, v_left, v_right)

    return _allocate(root, cov.columns, combine)

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)

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

    def combine(cluster: Cluster) -> Cluster:
        """Combine the child portfolios of a cluster via a Schur-augmented split."""
        left, right = _children(cluster)

        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 _split(cluster, v_left, v_right)

    return _allocate(root, cov.columns, combine)