"""Utils module for multivariate distribution."""
# Copyright (c) 2023-2026
# Author: Hugo Delatte <hugo.delatte@skfoliolabs.com>
# Credits: Matteo Manzi, Vincent Maladière, Carlo Nicolini
# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import auto
from functools import cached_property
from itertools import combinations
from typing import TYPE_CHECKING
import numpy as np
import scipy.sparse.csgraph as ssc
import scipy.stats as st
import sklearn.feature_selection as sf
from skfolio.typing import FloatArray
from skfolio.utils.tools import AutoEnum
if TYPE_CHECKING:
from skfolio.distribution.copula import BaseBivariateCopula
_MIXED_EDGE_ERROR = "An edge must connect two root nodes or two child nodes."
[docs]
class DependenceMethod(AutoEnum):
"""
Enumeration of methods to measure bivariate dependence.
Attributes
----------
KENDALL_TAU
Use Kendall's tau correlation coefficient.
MUTUAL_INFORMATION
Use mutual information estimated via a k-nearest neighbors method.
WASSERSTEIN_DISTANCE
Use the Wasserstein (Earth Mover's) distance.
"""
KENDALL_TAU = auto()
MUTUAL_INFORMATION = auto()
WASSERSTEIN_DISTANCE = auto()
@dataclass
class EdgeCondSets:
"""
Container for conditioning sets associated with an edge in an R-vine.
Attributes
----------
conditioned : tuple[int, int]
A tuple of conditioned variable indices.
conditioning : set[int]
A set of conditioning variable indices.
"""
conditioned: tuple[int, int]
conditioning: set[int]
def to_set(self) -> set[int]:
"""Union of conditioned and conditioning sets."""
return set(self.conditioned) | self.conditioning
def __add__(self, other: EdgeCondSets) -> EdgeCondSets:
"""Combine two EdgeCondSets, merging conditioned and conditioning sets."""
if not isinstance(other, self.__class__):
raise TypeError(
f"Cannot add a EdgeCondSets with an object of type {type(other)}"
)
s1 = self.to_set()
s2 = other.to_set()
conditioning = s1 & s2
first, second = s1 ^ s2
# maintain order
if first in other.conditioned:
first, second = second, first
return self.__class__(conditioned=(first, second), conditioning=conditioning)
def __repr__(self) -> str:
"""String representation of the EdgeCondSets."""
if self.conditioning:
return f"{self.conditioned} | {self.conditioning}"
return str(self.conditioned)
class BaseNode(ABC):
"""Base class for Nodes of the R-vine tree.
Attributes
----------
edges : set[Edge]
The set of edges attached to this node.
tree : Tree
The Tree containing this Node.
"""
def __init__(self) -> None:
self.edges: set[Edge] = set()
self.tree: Tree | None = None # Reference to the Tree containing this Node
@property
@abstractmethod
def ref(self) -> int | Edge:
"""Return the reference of this node (read-only).
For RootNode: reference of the variable index.
For ChildNode: reference of the edge in the previous tree.
"""
...
@abstractmethod
def clear_cache(self, clear_count: bool = True) -> None:
"""Clear the cached pseudo-values and margin values (u and v)."""
...
def __repr__(self) -> str:
"""String representation of the node."""
return f"Node({self.ref})"
class RootNode(BaseNode):
"""Root Node of the R-vine tree.
Parameters
----------
ref : int
The reference variable index.
central : bool
True if the node is central; otherwise, False.
pseudo_values : ndarray, optional
The pseudo-values of the Root Node.
Attributes
----------
edges : set[Edge]
The set of edges attached to this node.
tree : Tree
The Tree containing this Node.
"""
def __init__(
self, ref: int, central: bool, pseudo_values: FloatArray | None = None
) -> None:
super().__init__()
self._ref = ref
self.central = central
self.pseudo_values = pseudo_values
@property
def ref(self) -> int:
"""Return the variable index of this node (read-only)."""
return self._ref
def clear_cache(self, clear_count: bool = True) -> None:
"""Clear the cached pseudo-values."""
self.pseudo_values = None
class ChildNode(BaseNode):
"""Child Node of the R-vine tree.
A child node is an edge from the previous tree.
Parameters
----------
ref : Edge
The reference edge in the previous tree.
Attributes
----------
edges : set[Edge]
The set of edges attached to this node.
tree : Tree
The Tree containing this Node.
"""
def __init__(self, ref: Edge) -> None:
super().__init__()
self._ref = ref
# pointer from Edge to Node
ref.ref_node = self
self._central: bool | None = None
self._u: FloatArray | None = None
self._v: FloatArray | None = None
self._u_count: int = 0
self._v_count: int = 0
self._u_count_total: int = 0
self._v_count_total: int = 0
@property
def ref(self) -> Edge:
"""Return the edge of the previous tree referenced by this node (read-only)."""
return self._ref
@property
def central(self) -> bool:
"""Determine whether this node is considered central.
It is inherited from the associated edge's centrality.
Returns
-------
central: bool
True if the node is central; otherwise, False.
"""
if self._central is None:
self._central = self.ref.strongly_central
return self._central
@property
def u(self) -> FloatArray:
"""Get the first margin value (u) for the node.
It is obtained by computing the partial derivative of the copula with respect
to v.
Returns
-------
u : ndarray
The u values for this node.
"""
is_count = self.tree is not None and self.tree.is_count_visits
if is_count:
self._u_count_total += 1
else:
self._u_count += 1
if self._u is None:
X = self.ref.get_X()
if is_count:
self._u = np.array([np.nan])
else:
self._u = self.ref.copula.partial_derivative(X, first_margin=False)
value = self._u
# Clear cache
if (
not is_count
and self._u_count_total != 0
and self._u_count == self._u_count_total
):
self._u = None
self._u_count = 0
return value
@u.setter
def u(self, value: FloatArray) -> None:
"""Set the first margin value (u) for the node."""
self._u = value
@property
def v(self) -> FloatArray:
"""Get the second margin value (v) for the node.
It is obtained by computing the partial derivative of the copula with respect
to u.
Returns
-------
v : ndarray
The v values for this node.
"""
is_count = self.tree is not None and self.tree.is_count_visits
if is_count:
self._v_count_total += 1
else:
self._v_count += 1
if self._v is None:
X = self.ref.get_X()
if is_count:
self._v = np.array([np.nan])
else:
self._v = self.ref.copula.partial_derivative(X, first_margin=True)
value = self._v
# Clear cache
if (
not is_count
and self._v_count_total != 0
and self._v_count == self._v_count_total
):
self._v = None
self._v_count = 0
return value
@v.setter
def v(self, value: FloatArray) -> None:
"""Set the second margin value (v) for the node."""
self._v = value
def get_var(self, is_left: bool | None) -> int:
"""Return the variable index associated with this node.
The variable is determined by the conditioned set of the edge.
Parameters
----------
is_left : bool or None
Indicates whether to select the left or right node.
Returns
-------
var : int
The variable index corresponding to this node.
Raises
------
ValueError
If `is_left` is None.
"""
if is_left is None:
raise ValueError("is_left cannot be None for Child Nodes")
var = self.ref.cond_sets.conditioned[0 if is_left else 1]
return var
def clear_cache(self, clear_count: bool = True) -> None:
"""Clear the cached margin values (u and v) and counts.
Parameters
----------
clear_count : bool, default=True
If True, the visit counts are also reset.
"""
self._u = None
self._v = None
if clear_count:
self._u_count = 0
self._v_count = 0
self._u_count_total = 0
self._v_count_total = 0
class Edge:
"""
Represents an edge in an R-vine tree connecting two nodes.
This class encapsulates the information for an edge between two nodes in an R-vine,
including the associated copula, the dependence measure, and the conditioning sets.
Attributes
----------
node1 : RootNode | ChildNode
The first node in the edge.
node2 : RootNode | ChildNode
The second node in the edge.
dependence_method : DependenceMethod
The method used to measure dependence between the two nodes.
copula : BaseBivariateCopula
The fitted copula for this edge. Set when the vine copula is fitted.
ref_node : ChildNode
A pointer to the node in the next tree constructed from this edge. Set when
that node is created.
"""
copula: BaseBivariateCopula
ref_node: ChildNode
def __init__(
self,
node1: RootNode | ChildNode,
node2: RootNode | ChildNode,
dependence_method: DependenceMethod = DependenceMethod.KENDALL_TAU,
) -> None:
self.node1 = node1
self.node2 = node2
self.dependence_method = dependence_method
@cached_property
def weakly_central(self) -> bool:
"""Determine if the edge is weakly central.
An edge is weakly central if at least one of its two nodes is central.
"""
return self.node1.central or self.node2.central
@cached_property
def strongly_central(self) -> bool:
"""Determine if the edge is strongly central.
An edge is strongly central if both of its nodes are central.
"""
return self.node1.central and self.node2.central
@cached_property
def dependence(self) -> float:
"""Dependence measure between the two nodes.
This is computed on the data from the edge using the specified dependence
method.
"""
X = self.get_X()
dep = _dependence(X, dependence_method=self.dependence_method)
return dep
@cached_property
def cond_sets(self) -> EdgeCondSets:
"""Compute the conditioning sets for the edge.
For a root node edge, the conditioned set consists of the two variable indices.
For non-root nodes, the conditioning sets are obtained by combining the
conditioning sets of the two edges from the previous tree.
"""
node1, node2 = self.node1, self.node2
if isinstance(node1, RootNode) and isinstance(node2, RootNode):
return EdgeCondSets(conditioned=(node1.ref, node2.ref), conditioning=set())
if isinstance(node1, ChildNode) and isinstance(node2, ChildNode):
return node1.ref.cond_sets + node2.ref.cond_sets
raise TypeError(_MIXED_EDGE_ERROR)
def ref_to_nodes(self) -> None:
"""Connect this edge to its two nodes."""
self.node1.edges.add(self)
self.node2.edges.add(self)
def get_X(self) -> FloatArray:
"""Retrieve the bivariate pseudo-observation data associated with the edge.
For a root edge, this returns the pseudo-values from node1 and node2.
For non-root edges, the appropriate margins (u or v) are selected
based on the shared node order.
Returns
-------
X : ndarray of shape (n_observations, 2)
The bivariate pseudo-observation data corresponding to this edge.
"""
node1, node2 = self.node1, self.node2
if isinstance(node1, RootNode) and isinstance(node2, RootNode):
u = node1.pseudo_values
v = node2.pseudo_values
elif isinstance(node1, ChildNode) and isinstance(node2, ChildNode):
is_left1, is_left2 = node1.ref.shared_node_is_left(node2.ref)
u = node1.v if is_left1 else node1.u
v = node2.v if is_left2 else node2.u
else:
raise TypeError(_MIXED_EDGE_ERROR)
X = np.stack([u, v]).T
return X
def shared_node_is_left(self, other: Edge) -> tuple[bool, bool]:
"""Determine the ordering of shared nodes between this edge and another edge.
If the two edges share one node, this method indicates for each edge whether the
shared node is the left node.
Parameters
----------
other : Edge
Another edge to compare with.
Returns
-------
is_left1, is_left2 : tuple[bool, bool]
A tuple (is_left1, is_left2) where is_left1 is True if the shared node is
the left node of self and is_left2 is True if the shared node is the left
node of other.
Raises
------
ValueError
If the edges do not share exactly one node.
"""
if self.node1 == other.node1:
return True, True
if self.node2 == other.node1:
return False, True
if self.node2 == other.node2:
return False, False
# self.node1 == other.node2
raise ValueError("Edges are not correctly ordered")
def share_one_node(self, other: Edge) -> bool:
"""Check whether two edges share exactly one node.
Parameters
----------
other : Edge
Another edge to compare with.
Returns
-------
bool
True if the two edges share exactly one node; otherwise, False.
"""
return len({self.node1, self.node2} & {other.node1, other.node2}) == 1
def __repr__(self) -> str:
"""String representation of the edge."""
copula = getattr(self, "copula", None)
if copula is None:
return f"Edge({self.cond_sets})"
return f"Edge({self.cond_sets}, {copula.fitted_repr})"
class Tree:
"""
Represents an R-vine tree at level k.
A Tree consists of a set of nodes and the edges connecting them. It represents one
level (k) in the R-vine structure.
Parameters
----------
level : int
The tree level (k) in the R-vine.
nodes : list[Node]
A list of Node objects representing the nodes in this tree.
Attributes
----------
edges : list[Edge]
The list of edges in the Tree.
is_count_visits : bool
Whether to count the number of visit of each Node during sampling.
"""
def __init__(self, level: int, nodes: list[RootNode | ChildNode]) -> None:
self.level = level
self._nodes = nodes
for node in nodes:
# pointer from Node to Tree
node.tree = self
self.edges: list[Edge] = []
self.is_count_visits: bool = False
@property
def nodes(self) -> list[RootNode | ChildNode]:
"""Return the tree nodes (read-only)."""
return self._nodes
def set_edges_from_mst(self, dependence_method: DependenceMethod) -> None:
"""Construct the Maximum Spanning Tree (MST) from the current nodes using
the specified dependence method.
The MST is built based on pairwise dependence measures computed between nodes.
If any edge is (weakly) central, a central factor is added to the dependence
measure to favor edges connected to central nodes.
Parameters
----------
dependence_method : DependenceMethod
The method used to compute the dependence measure between nodes (e.g.,
Kendall's tau).
Returns
-------
None
"""
n = len(self.nodes)
dependence_matrix = np.zeros((n, n))
eligible_edges = {}
central = False
for i, j in combinations(range(n), 2):
node1 = self.nodes[i]
node2 = self.nodes[j]
# Above the first tree, two nodes can only be joined if their edges share
# one node (proximity condition).
if (
isinstance(node1, ChildNode)
and isinstance(node2, ChildNode)
and not node1.ref.share_one_node(node2.ref)
):
continue
edge = Edge(node1=node1, node2=node2, dependence_method=dependence_method)
if not central and edge.weakly_central:
central = True
# Negate the matrix to use minimum_spanning_tree for maximum spanning
# Add a cst to ensure that even if dep is 0, we still build a valid MST
dep = abs(edge.dependence) + 1e-5
dependence_matrix[i, j] = dep
eligible_edges[(i, j)] = edge
if np.any(np.isnan(dependence_matrix)):
raise RuntimeError("dependence_matrix contains NaNs")
if central:
max_dep = np.max(dependence_matrix)
for (i, j), edge in eligible_edges.items():
if edge.weakly_central:
if edge.strongly_central:
central_factor = 3 * max_dep
else:
central_factor = 2 * max_dep
dep = dependence_matrix[i, j] + central_factor
dependence_matrix[i, j] = dep
# Compute the minimum spanning tree
mst = ssc.minimum_spanning_tree(-dependence_matrix, overwrite=True)
edges = []
# Extract the indices of the non-zero entries (edges)
for i, j in zip(*mst.nonzero(), strict=True):
edge = eligible_edges[(i, j)]
# connect Nodes to Edges
edge.ref_to_nodes()
edges.append(edge)
self.edges = edges
def clear_cache(self, clear_count: bool = True) -> None:
"""Clear cached values for all nodes in the tree."""
for node in self.nodes:
node.clear_cache(clear_count=clear_count)
def __repr__(self) -> str:
"""String representation of the tree."""
return f"Tree(level {self.level})"
def _dependence(X: FloatArray, dependence_method: DependenceMethod) -> float:
"""Compute the dependence between two variables in X using the specified method.
Parameters
----------
X : array-like of shape (n_observations, 2)
A 2D array of bivariate inputs (u, v), where u and v are assumed to lie in
[0, 1].
dependence_method : DependenceMethod
The method to use for measuring dependence. Options are:
- DependenceMethod.KENDALL_TAU
- DependenceMethod.MUTUAL_INFORMATION
- DependenceMethod.WASSERSTEIN_DISTANCE
Returns
-------
dependence : float
The computed dependence measure.
Raises
------
ValueError
If X does not have exactly 2 columns or if an unsupported dependence method is
provided.
"""
X = np.asarray(X)
if X.ndim != 2 or X.shape[1] != 2:
raise ValueError("X must be a 2D array with exactly 2 columns.")
match dependence_method:
case DependenceMethod.KENDALL_TAU:
dep = st.kendalltau(X[:, 0], X[:, 1]).statistic
case DependenceMethod.MUTUAL_INFORMATION:
dep = sf.mutual_info_regression(X[:, 0].reshape(-1, 1), X[:, 1])[0]
case DependenceMethod.WASSERSTEIN_DISTANCE:
dep = st.wasserstein_distance(X[:, 0], X[:, 1])
case _:
raise ValueError(f"Dependence method {dependence_method} not valid")
return dep
