diff --git a/examples/plot_tree.py b/examples/plot_tree.py new file mode 100644 index 000000000..6f501215c --- /dev/null +++ b/examples/plot_tree.py @@ -0,0 +1,332 @@ +""" +======================================== +Tree wasserstein distance and barycenter +======================================== + +Illustrates the use of the tree wasserstein distance calculator, and barycenter solvers +for both the fixed-support and the free support versions +""" + +import matplotlib.pyplot as plt +import matplotlib.colors as colors +import ot +import torch + +import networkx as netx +import random +import matplotlib.pyplot as plt +import numpy as np + +import torch.nn.functional as F + +# Utilitaries functions + +norm = colors.SymLogNorm(linthresh=1e-3, vmin=0.0, vmax=1.0) + + +def tree_to_networkx(tree, length): + G = netx.Graph() + + nb_nodes = tree.shape[0] + + root = -1 + + for i in range(nb_nodes): + parent = int(tree[i]) + if i != parent: + G.add_edge(i, parent, weight=length[i]) + else: + root = i + + return G, root + + +# Source - https://stackoverflow.com/a/29597209 +# Posted by Joel, modified by community. See post 'Timeline' for change history +# Retrieved 2026-07-16, License - CC BY-SA 4.0 +def hierarchy_pos(G, root=None, width=1.0, vert_gap=0.2, vert_loc=0, xcenter=0.5): + """ + From Joel's answer at https://stackoverflow.com/a/29597209/2966723. + Licensed under Creative Commons Attribution-Share Alike + + If the graph is a tree this will return the positions to plot this in a + hierarchical layout. + + G: the graph (must be a tree) + + root: the root node of current branch + - if the tree is directed and this is not given, + the root will be found and used + - if the tree is directed and this is given, then + the positions will be just for the descendants of this node. + - if the tree is undirected and not given, + then a random choice will be used. + + width: horizontal space allocated for this branch - avoids overlap with other branches + + vert_gap: gap between levels of hierarchy + + vert_loc: vertical location of root + + xcenter: horizontal location of root + """ + if not netx.is_tree(G): + raise TypeError("cannot use hierarchy_pos on a graph that is not a tree") + + if root is None: + if isinstance(G, netx.DiGraph): + root = next( + iter(netx.topological_sort(G)) + ) # allows back compatibility with nx version 1.11 + else: + root = random.choice(list(G.nodes)) + + def _hierarchy_pos( + G, root, width=1.0, vert_gap=0.2, vert_loc=0, xcenter=0.5, pos=None, parent=None + ): + """ + see hierarchy_pos docstring for most arguments + + pos: a dict saying where all nodes go if they have been assigned + parent: parent of this branch. - only affects it if non-directed + + """ + + if pos is None: + pos = {root: (xcenter, vert_loc)} + else: + pos[root] = (xcenter, vert_loc) + children = list(G.neighbors(root)) + if not isinstance(G, netx.DiGraph) and parent is not None: + children.remove(parent) + if len(children) != 0: + dx = width / len(children) + nextx = xcenter - width / 2 - dx / 2 + for child in children: + nextx += dx + pos = _hierarchy_pos( + G, + child, + width=dx, + vert_gap=vert_gap, + vert_loc=vert_loc - vert_gap, + xcenter=nextx, + pos=pos, + parent=root, + ) + return pos + + return _hierarchy_pos(G, root, width, vert_gap, vert_loc, xcenter) + + +def print_tree(tree, length, measure=None, ax=None): + G, root = tree_to_networkx(tree, length) + + pos = hierarchy_pos(G, root=root) + + nb_nodes = tree.shape[0] + + if measure is None: + measure_color = [1 / nb_nodes for _ in range(nb_nodes)] + else: + measure_color = [measure[node] for node in G.nodes()] + + netx.draw_networkx_edges(G, pos=pos, ax=ax) + netx.draw_networkx_labels(G, pos=pos, ax=ax, font_color="white") + + measure_color = norm(measure_color) + + nodes = netx.draw_networkx_nodes( + G, + pos=pos, + ax=ax, + node_color=measure_color, + cmap=plt.cm.inferno, + node_size=300, + vmin=0.0, + vmax=1.0, + ) + + return nodes + + +# Gradient descent to transport a measure into an other +n = 7 +m = 7 + +s = ot.datasets.make_2D_samples_gauss(n, [0.0, 0.0], [[1.0, 0], [0, -1.0]], 0) +c = ot.datasets.make_2D_samples_gauss(m, [10.0, 10.0], [[1.0, 0], [0, -1.0]], 0) + +source_point = torch.tensor(s, dtype=torch.float32) +target_point = torch.tensor(c, dtype=torch.float32) + +points = torch.cat([source_point, target_point]) + +tree, length = ot.utils.random_tree(points, 30) + +source_mes = torch.zeros(n + m) +source_mes[:n] = torch.ones(n) / n + +source_mes_init = source_mes.clone() + +target_mes = torch.zeros(n + m) +target_mes[n:] = torch.ones(m) / m + +source_mes = source_mes.detach().requires_grad_(True) + +step = 0.001 + +for i in range(300): + loss = ot.lp.tree_wasserstein_distance(tree, length, source_mes, target_mes) + + loss.backward() + + with torch.no_grad(): + sourceNouv = source_mes - step * source_mes.grad + sourceNouv = ot.utils.proj_simplex(sourceNouv.cpu().numpy()) + + source_mes = torch.tensor(sourceNouv) + source_mes.requires_grad_(True) + +source_mes = source_mes.detach().numpy() + +fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 7)) + +print_tree(tree, length, source_mes_init, ax1) +ax1.set_title("Source measure") + +print_tree(tree, length, target_mes, ax2) +ax2.set_title("Target measure") + +last_nodes = print_tree(tree, length, source_mes, ax3) +ax3.set_title("Final measure") + +plt.suptitle( + "Gradient descent to minimize the TWD between source measure and target measure" +) + +real_values = [ + 0.0, + 0.001, + 0.01, + 0.02, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, +] + +transformed_values = norm(real_values) + +cbar = fig.colorbar( + last_nodes, + ax=[ax1, ax2, ax3], + orientation="horizontal", + shrink=0.7, + pad=0.03, + ticks=transformed_values, +) +cbar.ax.set_xticklabels([f"{v:.3f}" for v in real_values]) +cbar.set_label("Probability mass", labelpad=15) + +plt.show() + +# Fixed-support barycenter + +n = 10 +k = 5 +l = 3 + +tree, length = ot.utils.random_tree_fixed_leaves(n, k, np, 0) + +source_mes = np.random.rand(l) +source_mes = np.pad(source_mes, (0, k - l)) +source_mes /= np.sum(source_mes) + +target_mes = np.random.rand(k - l) +target_mes = np.pad(target_mes, (l, 0)) +target_mes /= np.sum(target_mes) + +barycenter = ot.lp.fixed_support_tree_barycenter( + tree, length, np.stack([source_mes, target_mes]) +) + +source_mes = np.pad(source_mes, (0, n - k)) +target_mes = np.pad(target_mes, (0, n - k)) +barycenter = np.pad(barycenter, (0, n - k)) + +fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 7)) + +print_tree(tree, length, source_mes, ax1) +ax1.set_title("Measure 1") + +print_tree(tree, length, target_mes, ax2) +ax2.set_title("Measure 2") + +last_nodes = print_tree(tree, length, barycenter, ax3) +ax3.set_title("Barycenter") + +plt.suptitle("Fixed support barycenter") + +cbar = fig.colorbar( + last_nodes, + ax=[ax1, ax2, ax3], + orientation="horizontal", + shrink=0.7, + pad=0.03, + ticks=transformed_values, +) +cbar.ax.set_xticklabels([f"{v:.3f}" for v in real_values]) +cbar.set_label("Probability mass", labelpad=15) + +plt.show() + +# Free support barycenter +n = 7 +m = 7 + +s = ot.datasets.make_2D_samples_gauss(n, [0.0, 0.0], [[1.0, 0], [0, -1.0]], 0) +c = ot.datasets.make_2D_samples_gauss(m, [10.0, 10.0], [[1.0, 0], [0, -1.0]], 0) + +source_point = torch.tensor(s, dtype=torch.float32) +target_point = torch.tensor(c, dtype=torch.float32) + +points = torch.cat([source_point, target_point]) + +tree, length = ot.utils.random_tree(points, 30) + +source_mes = torch.zeros(n + m) +target_mes = torch.zeros(n + m) +source_mes[12] = 1.0 +target_mes[11] = 1.0 + +barycenter = ot.lp.free_support_tree_barycenter( + tree, length, torch.stack([source_mes, target_mes]), 1000, 1e-1 +) + +fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 7)) + +print_tree(tree, length, source_mes, ax1) +ax1.set_title("Measure 1") + +print_tree(tree, length, target_mes, ax2) +ax2.set_title("Measure 2") + +last_nodes = print_tree(tree, length, barycenter, ax3) +ax3.set_title("Barycenter") + +plt.suptitle("Free support barycenter") + +cbar = fig.colorbar( + last_nodes, + ax=[ax1, ax2, ax3], + orientation="horizontal", + shrink=0.7, + pad=0.03, + ticks=transformed_values, +) +cbar.ax.set_xticklabels([f"{v:.3f}" for v in real_values]) +cbar.set_label("Probability mass", labelpad=15) + +plt.show() diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index c9fa676c4..1403b4e70 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -36,6 +36,16 @@ linear_circular_ot, ) +from .solver_tree import ( + topological_sort, + tree_wasserstein_distance, +) + +from .tree_barycenter import ( + fixed_support_tree_barycenter, + free_support_tree_barycenter, +) + __all__ = [ "emd", "emd2", @@ -60,4 +70,8 @@ "free_support_barycenter_generic_costs", "NorthWestMMGluing", "ot_barycenter_energy", + "topological_sort", + "tree_wasserstein_distance", + "fixed_support_tree_barycenter", + "free_support_tree_barycenter", ] diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py new file mode 100644 index 000000000..7624b0042 --- /dev/null +++ b/ot/lp/solver_tree.py @@ -0,0 +1,188 @@ +from ..backend import get_backend +import numpy as np +from collections import deque + +""" +Solver for the tree wasserstein distance +""" + +# Author : Ali Boudjema + + +def topological_sort(tree): + r""" + Computes a topological order of the given tree + + Parameters + ----------- + tree: array_like, shape(n) + ancestor of each node in the tree (ancestor of root is root) + """ + + n = tree.shape[0] + + in_degree = np.zeros(n, dtype=int) + + for cur_node in range(n): + if cur_node != tree[cur_node]: + in_degree[tree[cur_node]] += 1 + + queue = deque() + + for cur_node in range(n): + if in_degree[cur_node] == 0: + queue.append(cur_node) + + topo_order = [] + + while queue: + cur_node = queue.popleft() + topo_order.append(cur_node) + + ancestor = tree[cur_node] + + if cur_node != ancestor: + in_degree[ancestor] -= 1 + + if in_degree[ancestor] == 0: + queue.append(ancestor) + + return np.array(topo_order) + + +def tree_wasserstein_distance( + tree, length, u_weights, v_weights, topo_order=None, return_plans=False +): + r""" + Computes the tree wasserstein distance for a given tree between two empirical distributions + + Parameters + ---------- + tree : array_like, shape(n) + parent of each node in the tree (parent of root is root) + length : array_like, shape(n) + 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) + weights of the second empirical distributions + topo_order : array_like, shape(n), optional + topological order of the tree + return_plans : bool, optional + if True, returns the optimal transport plan between the + two distributions, default is False + + Returns + ------- + cost : float + The tree wasserstein distance + plans : coo_matrix, optional + If return_plans is True, returns a coo_matrix containing the plan + + Reference + --------- + The proof of this algorithm uses the formula (3) in the article + Tree-Sliced Variants of Wasserstein Distances + """ + + n = tree.shape[0] + + assert ( + n == length.shape[0] == u_weights.shape[0] == v_weights.shape[0] + ), "dimension error in the input" + + if topo_order is None: + topo_order = topological_sort(tree) + + nx = get_backend(tree, length, u_weights, v_weights) + + mass_dict = {} + + for cur in range(n): + if u_weights[cur] != v_weights[cur]: + mass_dict[cur] = {cur: u_weights[cur] - v_weights[cur]} + else: + mass_dict[cur] = {} + + source_plan = [] + sink_plan = [] + mass_plan = [] + + virt_size = [len(mass_dict[k]) for k in range(n)] + + cost = 0 + + depth = nx.zeros(n) + + for i in range(n - 2, -1, -1): + cur_node = topo_order[i] + depth[cur_node] = depth[tree[cur_node]] + length[cur_node] + + for cur in topo_order: + dict_cur = mass_dict[cur] + p = int(tree[cur]) + + if cur != p: + dict_p = mass_dict[p] + + if virt_size[cur] > virt_size[p]: + mass_dict[cur], mass_dict[p] = dict_p, dict_cur + dict_cur, dict_p = dict_p, dict_cur + virt_size[cur], virt_size[p] = virt_size[p], virt_size[cur] + + while len(dict_cur) > 0 and len(dict_p) > 0: + node_scur = next(iter(dict_cur)) + amount_scur = dict_cur[node_scur] + + node_sp = next(iter(dict_p)) + amount_sp = dict_p[node_sp] + + if (amount_scur > 0) != (amount_sp > 0): + match_amount = min(abs(amount_scur), abs(amount_sp)) + + source = node_scur if amount_scur > 0 else node_sp + sink = node_sp if amount_scur > 0 else node_scur + + source_plan.append(source) + sink_plan.append(sink) + mass_plan.append(match_amount) + + length_path = depth[source] + depth[sink] - 2 * depth[p] + cost = cost + match_amount * length_path + + if amount_scur > 0: + dict_cur[node_scur] = dict_cur[node_scur] - match_amount + dict_p[node_sp] = dict_p[node_sp] + match_amount + else: + dict_cur[node_scur] = dict_cur[node_scur] + match_amount + dict_p[node_sp] = dict_p[node_sp] - match_amount + + if dict_cur[node_scur] == 0: + del dict_cur[node_scur] + + if dict_p[node_sp] == 0: + del dict_p[node_sp] + + else: + dict_p[node_scur] = amount_scur + del dict_cur[node_scur] + + if len(dict_p) == 0: + mass_dict[cur], mass_dict[p] = dict_p, dict_cur + dict_cur, dict_p = dict_p, dict_cur + + virt_size[p] += virt_size[cur] + + if mass_plan: + mass_plan = nx.stack(mass_plan, axis=0) + source_plan = nx.from_numpy(np.asarray(source_plan), type_as=length) + sink_plan = nx.from_numpy(np.asarray(sink_plan), type_as=length) + + plans = nx.coo_matrix( + mass_plan, source_plan, sink_plan, shape=(n, n), type_as=length + ) + + if return_plans: + return cost, plans + else: + return cost diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py new file mode 100644 index 000000000..dad169aed --- /dev/null +++ b/ot/lp/tree_barycenter.py @@ -0,0 +1,270 @@ +from ..backend import get_backend +import numpy as np +from ..utils import proj_simplex +from ..utils import list_to_array +import torch + +from .solver_tree import topological_sort +from .solver_tree import tree_wasserstein_distance + +# Author : Ali Boudjema + + +def get_B_matrix(tree, length, nb_leafs): + nx = get_backend(length) + + rows = [] + col = [] + data = [] + + for cur_leaf in range(nb_leafs): + cur_node = cur_leaf + + while cur_node != tree[cur_node]: + rows.append(cur_node) + col.append(cur_leaf) + data.append(length[cur_node]) + cur_node = tree[cur_node] + + nb_rows = len(tree) + nb_col = nb_leafs + + data = list_to_array(data, nx=nx) + rows = list_to_array(rows, nx=nx) + col = list_to_array(col, nx=nx) + + B = nx.coo_matrix(data, rows, col, shape=(nb_rows, nb_col), type_as=length) + + return B + + +def get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes, nx): + idx = nx.zeros(nb_nodes, type_as=cur_B) + + for node in range(nb_nodes): + idx[node] = nx.searchsorted(B_mes_sorted[:, node], cur_B[node]) + 1 + + z = -nb_mes + 2 * idx - 2 + + g = nx.transpose(B) @ z + g /= nb_mes + + return g + + +def pre_process_trees(tree_list, length_list, measures): + nx = get_backend(length_list, measures) + + nb_leafs = measures.shape[2] + + prepared_trees = [] + + for tree, length, mes in zip(tree_list, length_list, measures): + B = get_B_matrix(tree, length, nb_leafs) + nb_mes = mes.shape[0] + + B_mes_list = [] + + for id_mes in range(nb_mes): + col = nx.reshape(mes[id_mes], (-1, 1)) + + res = B @ col + + res_dense = nx.todense(res) + + res_1d = nx.reshape(res_dense, (-1,)) + + B_mes_list.append(res_1d) + + B_mes = nx.stack(B_mes_list, axis=0) + + B_mes_sorted = nx.sort(B_mes, axis=0) + + prepared_trees.append( + { + "B": B, + "B_mes_sorted": B_mes_sorted, + "nb_nodes": tree.shape[0], + "nb_mes": nb_mes, + } + ) + + return prepared_trees + + +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 + A single tree of shape (n_t,) or an array 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 an array + 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_itr : 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(tree_list, length_list, measures) + + if tree_list.ndim == 1: + tree_list = np.reshape(tree_list, (1, *tree_list.shape)) + length_list = nx.reshape(length_list, (1, *length_list.shape)) + + if measures.ndim == 2: + measures = nx.reshape(measures, (1, *measures.shape)) + measures = nx.tile(measures, (tree_list.shape[0], 1, 1)) + + 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" + + prepared_trees = pre_process_trees(tree_list, length_list, measures) + + nb_leafs = measures.shape[2] + + if init_measure is None: + cur_mes = nx.ones(nb_leafs) / nb_leafs + else: + cur_mes = init_measure + + nb_tree = len(prepared_trees) + + for itr in range(nb_itr): + old_mes = nx.copy(cur_mes) + g_total = nx.zeros(nb_leafs) + + for tree_data in prepared_trees: + B = tree_data["B"] + B_mes_sorted = tree_data["B_mes_sorted"] + nb_nodes = tree_data["nb_nodes"] + nb_mes = tree_data["nb_mes"] + + cur_B = B @ cur_mes + + g_tree = get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes, nx) + g_total += g_tree + + g_mean = g_total / nb_tree + + g_mean /= nx.norm(g_mean) + 1e-12 + + cur_mes -= step * g_mean + + cur_mes = proj_simplex(cur_mes) + + if np.linalg.norm(cur_mes - old_mes) < tol: + break + + return cur_mes + + +def free_support_tree_barycenter( + tree, length, measures, nb_itr=100, step=0.01, weights=None +): + """Computes the tree wasserstein barycenter for a tree with no constraints on the + support of the measures + + Parameters + ---------- + tree : array_like + A tree of shape (n), the number of nodes. tree[i] contains the index + of the parent of node i (with tree[root] == root). + length : array_like + The edge weights corresponding to tree. A single array of shape (n_t) containing + the length of the edge connecting node i to its parent. + measures : array_like, shape (m,n) + The input probability distributions mapped to the nodes. + nb_itr : int, optional + the maximal number of iterations for the subgradient descent + step : float, optional + the step size of the descent + weights : array_like, shape(m), optional + The weight of each measure, set to uniform if none + """ + nb_nodes = tree.shape[0] + + barycenter = torch.ones(nb_nodes) / nb_nodes + topo_order = topological_sort(tree) + + tree = np.asarray(tree) + tree = torch.from_numpy(tree) + + length = np.asarray(length) + length = torch.from_numpy(length) + + measures = np.asarray(measures) + measures = torch.from_numpy(measures) + + nb_mes = measures.shape[0] + + if weights is None: + w = torch.ones(nb_mes) / nb_mes + else: + w = weights + + for itr in range(nb_itr): + barycenter.requires_grad_(True) + + loss = sum( + cur_w + * tree_wasserstein_distance(tree, length, cur_mes, barycenter, topo_order) + for cur_w, cur_mes in zip(w, measures) + ) + + loss.backward() + + grad = barycenter.grad + + if grad is None or torch.norm(grad) < 1e-12: + barycenter = barycenter.detach() + break + + with torch.no_grad(): + scaled_grad = step * grad + scaled_grad = scaled_grad - torch.max(scaled_grad) + + barycenter_next = barycenter * torch.exp(-scaled_grad) + + barycenter = barycenter_next / (torch.sum(barycenter_next) + 1e-15) + + barycenter = barycenter.detach() + + return barycenter diff --git a/ot/utils.py b/ot/utils.py index fcd823415..8dd387710 100644 --- a/ot/utils.py +++ b/ot/utils.py @@ -2148,3 +2148,77 @@ def split_sample_ratio( a2 = a02 return X_a1, X_a2, a1, a2, sel_a1, sel_a2 + + +def random_tree(points, seed=None): + """Generate a random tree structured from a set of geometric points. + + The tree is constructed by iteratively connecting each new node to a + randomly chosen existing node among the previous ones. + Edge lengths represent the Euclidean distance between connected points. + + Parameters + ---------- + points : array_like, shape (n, d) + Coordinates of the nodes in a d-dimensional space. The backend type + of this array (NumPy, PyTorch, etc.) determines the backend of the + output lengths. + + seed : int, optional + Seed for the random number generator to ensure reproducibility. + Default is None. + + Returns + ------- + tree : ndarray, shape (n,) + An integer array where each element represents the parent of that node. + By construction, tree[0] = 0 (the root is its own parent). + length : array_like, shape (n,) + Length of the edge above each node (distance to its parent). + The root has a length of 0. Match the backend type of `points`. + """ + nx = get_backend(points) + + if seed is not None: + nx.seed(seed) + np.random.seed(seed) + + nb_points = points.shape[0] + + tree = np.zeros(nb_points, dtype=int) + length = nx.full(nb_points, 0, type_as=points) + + for i in range(1, nb_points): + tree[i] = np.random.randint(0, i) + length[i] = nx.norm(points[i] - points[tree[i]]) + + tree = nx.from_numpy(tree) + + return tree, length + + +def random_tree_fixed_leaves(n, k, nx, seed=None): + """Generates a random tree with n nodes and k leaves, the leaves being the first nodes + in the tree, and nx is the specified backend + """ + + trash = nx.zeros(1) + + nx = get_backend(trash) + + if seed is not None: + np.random.seed(seed) + nx.seed(seed) + + tree = np.zeros(n, dtype=int) + length = nx.rand(n) + + length[n - 1] = 0 + + for i in range(n - 1): + tree[i] = np.random.randint(max(i + 1, n - k + 1), n) + tree[n - 1] = n - 1 + + tree = nx.from_numpy(tree) + + return tree, length diff --git a/test/test_tree_barycenter.py b/test/test_tree_barycenter.py new file mode 100644 index 000000000..fd94e2eee --- /dev/null +++ b/test/test_tree_barycenter.py @@ -0,0 +1,96 @@ +"""Tests for the tree wassesterin barycenter""" + +import numpy as np +import pytest + +import ot +from ot.lp import fixed_support_tree_barycenter +from ot.utils import list_to_array +from ot.lp import free_support_tree_barycenter + + +def test_fixed_symetry_reflexivity(nx): + n = 50 + k = 30 + + tree, length = ot.utils.random_tree_fixed_leaves(n, k, nx, 0) + + u = nx.rand(k) + u = u / nx.sum(u) + + v = nx.rand(k) + v = v / nx.sum(v) + + np.testing.assert_allclose( + fixed_support_tree_barycenter(tree, length, nx.stack([u, v])), + fixed_support_tree_barycenter(tree, length, nx.stack([v, u])), + rtol=1e-2, + atol=1e-3, + ) + + np.testing.assert_allclose( + fixed_support_tree_barycenter( + tree, length, nx.stack([u, u]), nb_itr=100, step=0.01 + ), + u, + rtol=5.0, + atol=1e-3, + ) + + +def test_fixed_multiple_trees(nx): + n = 40 + k = 20 + + nb_trees = 5 + + data = [ot.utils.random_tree_fixed_leaves(n, k, nx, 41) for _ in range(nb_trees)] + trees, lengths = map(list, zip(*data)) + + trees = list_to_array(trees, nx=nx) + lengths = list_to_array(lengths, nx=nx) + + u = nx.rand(k) + u = u / nx.sum(u) + + np.testing.assert_allclose( + fixed_support_tree_barycenter( + trees, lengths, nx.stack([u, u]), nb_itr=1000, step=0.001 + ), + u, + rtol=1e-3, + atol=1e-3, + ) + + +def test_free_reflexivity_symetry(nx): + n = 10 + mu = np.array([0.0, 0.0]) + sigma = np.array([[1.0, 0.0], [0.0, 1.0]]) + points_np = ot.datasets.make_2D_samples_gauss(n, mu, sigma) + points = nx.from_numpy(points_np) + + tree, length = ot.utils.random_tree(points, 0) + length *= 100 + + u = nx.rand(n) + u = u / nx.sum(u) + + v = nx.rand(n) + v = v / nx.sum(v) + + np.testing.assert_allclose( + free_support_tree_barycenter(tree, length, nx.stack([u, v])), + free_support_tree_barycenter(tree, length, nx.stack([v, u])), + rtol=1e-2, + atol=1e-3, + ) + + np.testing.assert_allclose( + free_support_tree_barycenter( + tree, length, nx.stack([u, u]), nb_itr=100, step=0.0005 + ), + u, + rtol=5.0, + atol=1e-3, + ) diff --git a/test/test_tree_distance.py b/test/test_tree_distance.py new file mode 100644 index 000000000..5b4bf598b --- /dev/null +++ b/test/test_tree_distance.py @@ -0,0 +1,61 @@ +"""Tests for the tree wassesterin distance""" + +import pytest +import numpy as np + +import ot +from ot.lp import tree_wasserstein_distance +from ot.lp import wasserstein_1d + + +def test_symetry_reflexivity(nx): + n = 50 + mu = np.array([0.0, 0.0]) + sigma = np.array([[1.0, 0.0], [0.0, 1.0]]) + points_np = ot.datasets.make_2D_samples_gauss(n, mu, sigma) + points = nx.from_numpy(points_np) + + tree, length = ot.utils.random_tree(points) + + u = nx.rand(n) + u = u / nx.sum(u) + + v = nx.rand(n) + v = v / nx.sum(v) + + np.testing.assert_almost_equal( + tree_wasserstein_distance(tree, length, u, v), + tree_wasserstein_distance(tree, length, v, u), + ) + + np.testing.assert_almost_equal(tree_wasserstein_distance(tree, length, u, u), 0) + + +def test_chain(nx): + n = 50 + mu = 0 + sigma = 1 + + points_np = np.random.normal(mu, sigma, n) + points_np = np.sort(points_np) + points = nx.from_numpy(points_np) + + tree = nx.arange(n - 1, -1) + tree[0] = 0 + + length = nx.zeros(n) + + for i in range(1, n): + length[i] = points[i] - points[i - 1] + + u = nx.rand(n) + u = u / nx.sum(u) + + v = nx.rand(n) + v = v / nx.sum(v) + + np.testing.assert_almost_equal( + tree_wasserstein_distance(tree, length, u, v), + wasserstein_1d(points, points, u, v), + decimal=4, + )