doc

Author: Boudjema Ali

ot/lp/__init__.py

@@ -38,13 +38,11 @@
 
 from .solver_tree import (
     topological_sort,
-    tree_wasserstein,
+    tree_wasserstein_distance,
 )
 
 from .tree_barycenter import (
-    tree_barycenter,
     fixed_support_tree_barycenter,
-    sliced_fixed_support_tree_barycenter,
 )
 
 __all__ = [
@@ -72,8 +70,6 @@
     "NorthWestMMGluing",
     "ot_barycenter_energy",
     "topological_sort",
-    "tree_wasserstein",
-    "tree_barycenter",
+    "tree_wasserstein_distance",
     "fixed_support_tree_barycenter",
-    "sliced_fixed_support_tree_barycenter",
 ]

ot/lp/solver_tree.py

@@ -50,7 +50,7 @@ def topological_sort(tree):
     return np.array(topo_order)
 
 
-def tree_wasserstein(
+def tree_wasserstein_distance(
     tree, length, u_weights, v_weights, topo_order=None, return_plans=False
 ):
     r"""
@@ -59,9 +59,9 @@ def tree_wasserstein(
     Parameters
     ----------
     tree : array_like, shape(n)
-        ancestor of each node in the tree (ancestor of root is root)
+        parent of each node in the tree (parent of root is root)
     length : array_like, shape(n)
-        length of the arc above each node (length of root is 0)
+        length of the edge above each node (length of root is 0)
     u_weights : array_like, shape(n)
         weights of the first empirical distributions
     v_weights : array_like, shape(n)

ot/lp/tree_barycenter.py

@@ -1,6 +1,5 @@
 from ..backend import get_backend
 import numpy as np
-from .solver_tree import topological_sort
 from ..utils import proj_simplex
 from ..utils import list_to_array
 
@@ -9,101 +8,6 @@
 # IMPORTANT : ON PREND COMME CONVENTION QUE LES FEUILLES SONT LES PREMIERS SOMMETS DE L'ARBRE
 
 
-def wgm(values, weights):
-    # Returns the weighted geometric median
-
-    nx = get_backend(values, weights)
-
-    sorted_indices = np.argsort(values, kind="stable")
-
-    values_sorted = values[sorted_indices]
-    weights_sorted = weights[sorted_indices]
-
-    cum_weights = nx.cumsum(weights_sorted)
-
-    id = nx.searchsorted(cum_weights, 0.5 - 1e9)
-
-    return values_sorted[id]
-
-
-def get_measure(z, tree, length):
-    # Retrieves the measure from a vector after the wgm
-
-    n = z.shape[0]
-
-    nx = get_backend(length)
-
-    measure = nx.zeros(n)
-
-    for i in range(n):
-        p = tree[i]
-
-        if i == p:
-            measure[i] += 1
-        else:
-            measure[i] += z[i] / length[i]
-            measure[p] -= z[i] / length[i]
-
-    return measure
-
-
-def tree_barycenter(tree, length, measure, weights, topo_order=None):
-    r"""
-    Computes the tree wasserstein barycenter for a given tree between multiplie empirical distributions
-
-    Parameters
-    ----------
-    tree : array_like, shape(n)
-        ancestor of each node in the tree (ancestor of root is root)
-    length : array_like, shape(n)
-        length of the arc above each node (length of root is 0)
-    measure : array_like, shape(m, n)
-        distributions in the tree
-    weights : array_like, shape(m)
-        weight of each distribution
-
-    Returns
-    -------
-    barycenter : array_like, shape(n)
-        distribution of the barycenter
-
-    Reference
-    ---------
-    The code is a direct implementation of the algorithm described in
-    Tree-Wasserstein Barycenter for Large-Scale Multilevel Clustering and Scalable Bayes
-
-    """
-    n_measure = measure.shape[0]
-    n_node = tree.shape[0]
-
-    assert n_measure == weights.shape[0], "dimension error"
-
-    nx = get_backend(measure, weights, length)
-
-    z_measure = nx.zeros((n_measure, n_node))
-
-    if topo_order is None:
-        topo_order = topological_sort(tree)
-
-    for cur_node in topo_order:
-        p = tree[cur_node]
-
-        for id_mes in range(n_measure):
-            z_measure[id_mes][cur_node] += measure[id_mes][cur_node]
-
-            if cur_node != p:
-                z_measure[id_mes][p] += z_measure[id_mes][cur_node]
-
-    z = nx.zeros(n_node)
-
-    for cur_node in range(n_node):
-        z_measure[:, cur_node] *= length[cur_node]
-
-        z[cur_node] = wgm(z_measure[:, cur_node], weights)
-
-    return get_measure(z, tree, length)
-
-
 def get_B_matrix(tree, length, nb_leafs):
     nx = get_backend(length)
 
@@ -141,38 +45,14 @@ def get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes):
     return g
 
 
-def fixed_support_tree_barycenter(tree, length, measures, nb_itr=100, step=0.1):
-    nx = get_backend(length, measures)
-    nb_leafs = measures.shape[1]
-
-    B = get_B_matrix(tree, length, nb_leafs)
-
-    nb_mes = measures.shape[0]
-    nb_nodes = tree.shape[0]
-
-    cur_mes = nx.ones(nb_leafs) / nb_leafs
-
-    B_mes = list_to_array([B.dot(measures[i]) for i in range(nb_mes)])
-
-    sigma = nx.argsort(B_mes, axis=0)
-
-    B_mes_sorted = nx.take_along_axis(B_mes, sigma, axis=0)
-
-    for itr in range(nb_itr):
-        cur_B = B.dot(cur_mes)
-
-        g = get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes)
-
-        cur_mes -= step * g
-
-        cur_mes = proj_simplex(cur_mes)
-
-    return cur_mes
-
-
 def pre_process_trees(tree_list, length_list, measures):
     nx = get_backend(length_list, measures)
 
+    if tree_list.ndim == 1:
+        tree_list = nx.reshape(tree_list, (1, *tree_list.shape))
+        length_list = nx.reshape(length_list, (1, *length_list.shape))
+        measures = nx.reshape(measures, (1, *measures.shape))
+
     nb_leafs = measures.shape[2]
 
     prepared_trees = []
@@ -197,22 +77,69 @@ def pre_process_trees(tree_list, length_list, measures):
     return prepared_trees
 
 
-def sliced_fixed_support_tree_barycenter(
-    tree_list, length_list, measures, nb_itr=100, step=0.01, tol=1e-5
+def fixed_support_tree_barycenter(
+    tree_list, length_list, measures, nb_itr=100, step=0.01, tol=1e-5, init_measure=None
 ):
     """
+    Computes the Tree-Wasserstein (or Tree-Sliced) barycenter for one or multiple trees,
+    with the constraint that the support of the barycenter is fixed at the leaves.
+    It is assumed that the leaves correspond to the first nodes of the tree (indices 0 to k-1).
+    While the number of leaves (k) must be strictly identical across all structures to ensure
+    consistent alignment, the total number of nodes (leaves + internal nodes) can
+    freely vary from one tree to another.
+
+    If a single tree structure is provided (e.g., 1D arrays for tree_list/length_list
+    and 2D for measures), the function automatically expands their dimensions to 3D
+    internal structures to handle them uniformly as a multi-tree setting with t=1.
+
     Parameters
     -----------
-    tree_list : array_like, shape (t, n)
-    length_list : array_like, shape (t, n)
-    measures : array_like, shape (t, m, k)
+    tree_list : array_like
+        A single tree of shape (n_t,) or a list of t trees where each tree has shape (n_t,).
+        n_t is the total number of nodes in that specific tree. tree[i] contains the index
+        of the parent of node i (with tree[root] == root).
+    length_list : array_like
+        The edge weights corresponding to tree_list. A single array of shape (n_t,) or a list
+        of t arrays, where the t-th array has shape (n_t,) and contains the length of the edge
+        connecting node i to its parent.
+    measures : array_like, shape (t, m, k) or (m,k)
+        The input probability distributions mapped to the leaves.
+        k is the fixed number of leaves shared by all trees, and m is the number
+        of measures. Accepts a 2D array of shape (m, k) for a single tree, or a 3D array
+        of shape (t, m, k) in a multi-tree setting.
+    nb_tr : int, optional
+        the maximal number of iterations for the subgradient descent
+    step : float, optional
+        the step size of the descent
+    tol : float, optional
+        Convergence tolerance. The descent stops if the L2 norm of the difference
+        between two consecutive iterations is smaller than tol.
+    init_measure : array_like, shape (k), optional
+        The starting point of the descent, default is None
+
+    Returns
+    -------
+    cur_mes : array_like, shape (k)
+        The computed fixed-support barycenter supported on the k leaves.
+
+    References
+    ----------
+    "Fixed Support Tree-Sliced Wasserstein Barycenter"
     """
 
     nx = get_backend(length_list, measures)
 
+    assert (
+        tree_list.shape[0] == length_list.shape[0] == measures.shape[0]
+        and tree_list.shape[1] == length_list.shape[1]
+    ), "dimension error in the input"
+
     nb_leafs = measures.shape[2]
 
-    cur_mes = nx.ones(nb_leafs) / nb_leafs
+    if init_measure is None:
+        cur_mes = nx.ones(nb_leafs) / nb_leafs
+    else:
+        cur_mes = init_measure
 
     prepared_trees = pre_process_trees(tree_list, length_list, measures)