From 08d417e4de98f7ec3bf06af4bfc367b66baedc40 Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Wed, 10 Jun 2026 15:29:29 +0200 Subject: [PATCH 01/13] twd --- ot/lp/__init__.py | 7 +++ ot/lp/solver_tree.py | 103 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 ot/lp/solver_tree.py diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index c9fa676c4..c42a0fefd 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -36,6 +36,11 @@ linear_circular_ot, ) +from .solver_tree import ( + topological_sort, + tree_wasserstein, +) + __all__ = [ "emd", "emd2", @@ -60,4 +65,6 @@ "free_support_barycenter_generic_costs", "NorthWestMMGluing", "ot_barycenter_energy", + "topological_sort", + "tree_wasserstein", ] diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py new file mode 100644 index 000000000..243a37b01 --- /dev/null +++ b/ot/lp/solver_tree.py @@ -0,0 +1,103 @@ +from ..backend import get_backend +import numpy as np +from collections import deque + +""" +Solver for the tree wasserstein distance problem +""" + +# Author : Ali Boudjema + + +def topological_sort(tree): + r""" + Computes a topological sort of the given tree + The tree is an array : tree[i] is the direct ancestor of node i, and tree[root] = root + + Je ne vérifie pas que l'arbre est du bon format + Pas de backend ici ? + """ + + n = len(tree) + + in_degree = np.zeros(n) + + 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) + + if cur_node != tree[cur_node]: + in_degree[tree[cur_node]] -= 1 + + if in_degree[tree[cur_node]] == 0: + queue.append(tree[cur_node]) + + return np.array(topo_order) + + +def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): + r""" + Computes the tree wasserstein distance for a given tree between two 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) + 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) + topological order of the tree, optional + + Returns + ------- + cost : float/array_like, shape(...) + The tree wasserstein distance + """ + + n = len(tree) + + assert ( + n == len(length) == 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(length, u_weights, v_weights) + + u_cumweights = nx.copy(u_weights) + v_cumweights = nx.copy(v_weights) + + cost = 0 + + for i in range(n): + cur_node = topo_order[i] + + cost += length[cur_node] * nx.abs( + u_cumweights[cur_node] - v_cumweights[cur_node] + ) + + u_cumweights[tree[cur_node]] = ( + u_cumweights[cur_node] + u_cumweights[tree[cur_node]] + ) + v_cumweights[tree[cur_node]] = ( + v_cumweights[cur_node] + v_cumweights[tree[cur_node]] + ) + + return cost From 8dd3ba5c7e08880737bb133e19b3f845a3dc2d4e Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Thu, 11 Jun 2026 09:56:56 +0200 Subject: [PATCH 02/13] twd --- ot/lp/solver_tree.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py index 243a37b01..781a27480 100644 --- a/ot/lp/solver_tree.py +++ b/ot/lp/solver_tree.py @@ -11,14 +11,15 @@ def topological_sort(tree): r""" - Computes a topological sort of the given tree - The tree is an array : tree[i] is the direct ancestor of node i, and tree[root] = root + Computes a topological order of the given tree - Je ne vérifie pas que l'arbre est du bon format - Pas de backend ici ? + Parameters + ----------- + tree: array_like, shape(n, ...) + ancestor of each node in the tree (ancestor of root is root) """ - n = len(tree) + n = tree.shape[0] in_degree = np.zeros(n) @@ -53,9 +54,9 @@ def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): Parameters ---------- - tree : array_like, shape(n) + tree : array_like, shape(n, ...) ancestor of each node in the tree (ancestor of root is root) - length : array_like, shape(n) + length : array_like, shape(n, ...) length of the arc above each node (length of root is 0) u_weights : array_like, shape(n, ...) weights of the first empirical distributions @@ -70,10 +71,10 @@ def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): The tree wasserstein distance """ - n = len(tree) + n = tree.shape[0] assert ( - n == len(length) == u_weights.shape[0] == v_weights.shape[0] + n == length.shape[0] == u_weights.shape[0] == v_weights.shape[0] ), "dimension error in the input" if topo_order is None: From 2e9ac447fcb9a5ec26db3c4fcfd9868dbb1e8df1 Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Mon, 15 Jun 2026 15:43:49 +0200 Subject: [PATCH 03/13] twd --- ot/lp/solver_tree.py | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py index 781a27480..d9a2e59be 100644 --- a/ot/lp/solver_tree.py +++ b/ot/lp/solver_tree.py @@ -3,7 +3,7 @@ from collections import deque """ -Solver for the tree wasserstein distance problem +Solver for the tree wasserstein distance """ # Author : Ali Boudjema @@ -15,13 +15,13 @@ def topological_sort(tree): Parameters ----------- - tree: array_like, shape(n, ...) + 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) + in_degree = np.zeros(n, dtype=int) for cur_node in range(n): if cur_node != tree[cur_node]: @@ -39,11 +39,13 @@ def topological_sort(tree): cur_node = queue.popleft() topo_order.append(cur_node) - if cur_node != tree[cur_node]: - in_degree[tree[cur_node]] -= 1 + ancestor = tree[cur_node] + + if cur_node != ancestor: + in_degree[ancestor] -= 1 - if in_degree[tree[cur_node]] == 0: - queue.append(tree[cur_node]) + if in_degree[ancestor] == 0: + queue.append(ancestor) return np.array(topo_order) @@ -54,20 +56,20 @@ def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): Parameters ---------- - tree : array_like, shape(n, ...) + tree : array_like, shape(n) ancestor of each node in the tree (ancestor of root is root) - length : array_like, shape(n, ...) + length : array_like, shape(n) length of the arc above each node (length of root is 0) - u_weights : array_like, shape(n, ...) + u_weights : array_like, shape(n) weights of the first empirical distributions - v_weights : array_like, shape(n, ...) + v_weights : array_like, shape(n) weights of the second empirical distributions topo_order : array_like, shape(n) topological order of the tree, optional Returns ------- - cost : float/array_like, shape(...) + cost : float The tree wasserstein distance """ @@ -94,11 +96,9 @@ def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): u_cumweights[cur_node] - v_cumweights[cur_node] ) - u_cumweights[tree[cur_node]] = ( - u_cumweights[cur_node] + u_cumweights[tree[cur_node]] - ) - v_cumweights[tree[cur_node]] = ( - v_cumweights[cur_node] + v_cumweights[tree[cur_node]] - ) + ancestor = tree[cur_node] + + u_cumweights[ancestor] = u_cumweights[cur_node] + u_cumweights[ancestor] + v_cumweights[ancestor] = v_cumweights[cur_node] + v_cumweights[ancestor] return cost From 7c9425b63ac9f87c8aa4144b25983a48be56ee59 Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Tue, 16 Jun 2026 15:53:34 +0200 Subject: [PATCH 04/13] plan transport --- ot/lp/solver_tree.py | 100 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 87 insertions(+), 13 deletions(-) diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py index d9a2e59be..92a470826 100644 --- a/ot/lp/solver_tree.py +++ b/ot/lp/solver_tree.py @@ -50,7 +50,9 @@ def topological_sort(tree): return np.array(topo_order) -def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): +def tree_wasserstein( + 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 @@ -64,13 +66,18 @@ def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): weights of the first empirical distributions v_weights : array_like, shape(n) weights of the second empirical distributions - topo_order : array_like, shape(n) - topological order of the tree, optional + 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 """ n = tree.shape[0] @@ -84,21 +91,88 @@ def tree_wasserstein(tree, length, u_weights, v_weights, topo_order=None): nx = get_backend(length, u_weights, v_weights) - u_cumweights = nx.copy(u_weights) - v_cumweights = nx.copy(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 - for i in range(n): + 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] - cost += length[cur_node] * nx.abs( - u_cumweights[cur_node] - v_cumweights[cur_node] - ) + for cur in topo_order: + dict_cur = mass_dict[cur] + p = tree[cur] - ancestor = tree[cur_node] + 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 += match_amount * length_path + + if amount_scur > 0: + dict_cur[node_scur] -= match_amount + dict_p[node_sp] += match_amount + else: + dict_cur[node_scur] += match_amount + 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] - u_cumweights[ancestor] = u_cumweights[cur_node] + u_cumweights[ancestor] - v_cumweights[ancestor] = v_cumweights[cur_node] + v_cumweights[ancestor] + plans = nx.coo_matrix( + mass_plan, source_plan, sink_plan, shape=(n, n), type_as=length + ) - return cost + if return_plans: + return cost, plans + else: + return cost From bdd3f9e66c8f9b5c480206f7d3d128288a9b222f Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Wed, 17 Jun 2026 15:02:40 +0200 Subject: [PATCH 05/13] barycenter --- ot/lp/__init__.py | 5 +++ ot/lp/solver_tree.py | 5 +++ ot/lp/tree_barycenter.py | 93 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 103 insertions(+) create mode 100644 ot/lp/tree_barycenter.py diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index c42a0fefd..99c9c5cd5 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -41,6 +41,10 @@ tree_wasserstein, ) +from .tree_barycenter import ( + tree_barycenter, +) + __all__ = [ "emd", "emd2", @@ -67,4 +71,5 @@ "ot_barycenter_energy", "topological_sort", "tree_wasserstein", + "tree_barycenter", ] diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py index 92a470826..498d4778c 100644 --- a/ot/lp/solver_tree.py +++ b/ot/lp/solver_tree.py @@ -78,6 +78,11 @@ def tree_wasserstein( 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] diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py new file mode 100644 index 000000000..046559aae --- /dev/null +++ b/ot/lp/tree_barycenter.py @@ -0,0 +1,93 @@ +from ..backend import get_backend + + +def wgm(values, weights): + # Returns the weighted geometric median + + nx = get_backend(values, weights) + + sorted_indices = nx.argsort(values) + + values_sorted = values[sorted_indices] + weights_sorted = weights[sorted_indices] + + cum_weights = nx.cumsum(weights_sorted) + + id = nx.searchsorted(cum_weights, 0.5) + + 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): + 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)) + + for cur_node in range(n_node): + 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] += 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) From 97c1e1b21315bbc538ee4d42c3728c767949cbde Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Fri, 19 Jun 2026 11:12:04 +0200 Subject: [PATCH 06/13] correction bug barycenter --- ot/lp/tree_barycenter.py | 59 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 54 insertions(+), 5 deletions(-) diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py index 046559aae..c25ea5542 100644 --- a/ot/lp/tree_barycenter.py +++ b/ot/lp/tree_barycenter.py @@ -1,4 +1,50 @@ from ..backend import get_backend +import numpy as np +from collections import deque + +# Author : Ali Boudjema + + +# A retirer +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 wgm(values, weights): @@ -6,14 +52,14 @@ def wgm(values, weights): nx = get_backend(values, weights) - sorted_indices = nx.argsort(values) + 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) + id = nx.searchsorted(cum_weights, 0.5 - 1e9) return values_sorted[id] @@ -39,7 +85,7 @@ def get_measure(z, tree, length): return measure -def tree_barycenter(tree, length, measure, weights): +def tree_barycenter(tree, length, measure, weights, topo_order=None): r""" Computes the tree wasserstein barycenter for a given tree between multiplie empirical distributions @@ -74,14 +120,17 @@ def tree_barycenter(tree, length, measure, weights): z_measure = nx.zeros((n_measure, n_node)) - for cur_node in range(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] += measure[id_mes][cur_node] + z_measure[id_mes][p] += z_measure[id_mes][cur_node] z = nx.zeros(n_node) From f7d822a6659910012f14a043782e6a25ea992ea6 Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Fri, 19 Jun 2026 11:37:05 +0200 Subject: [PATCH 07/13] tree barycenter --- ot/lp/tree_barycenter.py | 44 +--------------------------------------- 1 file changed, 1 insertion(+), 43 deletions(-) diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py index c25ea5542..38f1e49e1 100644 --- a/ot/lp/tree_barycenter.py +++ b/ot/lp/tree_barycenter.py @@ -1,52 +1,10 @@ from ..backend import get_backend import numpy as np -from collections import deque +from .solver_tree import topological_sort # Author : Ali Boudjema -# A retirer -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 wgm(values, weights): # Returns the weighted geometric median From aadd8b11fd3ad964420cfe0623a57f617fddf6d1 Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Tue, 30 Jun 2026 15:14:01 +0200 Subject: [PATCH 08/13] tree barycenter gradient descent --- ot/lp/__init__.py | 2 ++ ot/lp/tree_barycenter.py | 65 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index 99c9c5cd5..68ee4faba 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -43,6 +43,7 @@ from .tree_barycenter import ( tree_barycenter, + fixed_support_tree_barycenter, ) __all__ = [ @@ -72,4 +73,5 @@ "topological_sort", "tree_wasserstein", "tree_barycenter", + "fixed_support_tree_barycenter", ] diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py index 38f1e49e1..95e68e0f2 100644 --- a/ot/lp/tree_barycenter.py +++ b/ot/lp/tree_barycenter.py @@ -1,9 +1,12 @@ from ..backend import get_backend import numpy as np from .solver_tree import topological_sort +from ..utils import proj_simplex # Author : Ali Boudjema +# IMPORTANT : ON PREND COMME CONVENTION QUE LES FEUILLES SONT LES PREMIERS SOMMETS DE L'ARBRE + def wgm(values, weights): # Returns the weighted geometric median @@ -98,3 +101,65 @@ def tree_barycenter(tree, length, measure, weights, topo_order=None): 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) + + 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 + + B = nx.coo_matrix(data, rows, col, shape=(nb_rows, nb_col), type_as=length) + + return B + + +def fixed_support_tree_barycenter(tree, length, measures, nb_itr=100, step=0.1): + nx = get_backend(length) + 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 = [B.dot(measures[i]) for i in range(nb_mes)] + + B_mes = np.asarray(B_mes) + + sigma = np.argsort(B_mes, axis=0) + + B_mes_sorted = np.take_along_axis(B_mes, sigma, axis=0) + + for itr in range(nb_itr): + cur_B = B.dot(cur_mes) + + idx = np.zeros(nb_nodes, dtype=int) + + for node in range(nb_nodes): + idx[node] = np.searchsorted(B_mes_sorted[:, node], cur_B[node]) + 1 + + z = -nb_mes + 2 * idx - 2 + + g = B.T.dot(z) / nb_mes + + cur_mes -= step * g + + cur_mes = proj_simplex(cur_mes) + + return cur_mes From b042821a9702a82855eb289ba021886bab97f2f8 Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Wed, 1 Jul 2026 15:09:04 +0200 Subject: [PATCH 09/13] tree sliced barycenter --- ot/lp/__init__.py | 2 + ot/lp/tree_barycenter.py | 106 ++++++++++++++++++++++++++++++++++----- 2 files changed, 96 insertions(+), 12 deletions(-) diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index 68ee4faba..3aa010a1f 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -44,6 +44,7 @@ from .tree_barycenter import ( tree_barycenter, fixed_support_tree_barycenter, + sliced_fixed_support_tree_barycenter, ) __all__ = [ @@ -74,4 +75,5 @@ "tree_wasserstein", "tree_barycenter", "fixed_support_tree_barycenter", + "sliced_fixed_support_tree_barycenter", ] diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py index 95e68e0f2..f8b58c19f 100644 --- a/ot/lp/tree_barycenter.py +++ b/ot/lp/tree_barycenter.py @@ -2,6 +2,7 @@ import numpy as np from .solver_tree import topological_sort from ..utils import proj_simplex +from ..utils import list_to_array # Author : Ali Boudjema @@ -127,8 +128,21 @@ def get_B_matrix(tree, length, nb_leafs): return B +def get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes): + idx = np.zeros(nb_nodes, dtype=int) + + for node in range(nb_nodes): + idx[node] = np.searchsorted(B_mes_sorted[:, node], cur_B[node]) + 1 + + z = -nb_mes + 2 * idx - 2 + + g = B.T.dot(z) / nb_mes + + return g + + def fixed_support_tree_barycenter(tree, length, measures, nb_itr=100, step=0.1): - nx = get_backend(length) + nx = get_backend(length, measures) nb_leafs = measures.shape[1] B = get_B_matrix(tree, length, nb_leafs) @@ -138,28 +152,96 @@ def fixed_support_tree_barycenter(tree, length, measures, nb_itr=100, step=0.1): cur_mes = nx.ones(nb_leafs) / nb_leafs - B_mes = [B.dot(measures[i]) for i in range(nb_mes)] - - B_mes = np.asarray(B_mes) + B_mes = list_to_array([B.dot(measures[i]) for i in range(nb_mes)]) - sigma = np.argsort(B_mes, axis=0) + sigma = nx.argsort(B_mes, axis=0) - B_mes_sorted = np.take_along_axis(B_mes, sigma, 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) - idx = np.zeros(nb_nodes, dtype=int) + g = get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes) + + cur_mes -= step * g + + cur_mes = proj_simplex(cur_mes) - for node in range(nb_nodes): - idx[node] = np.searchsorted(B_mes_sorted[:, node], cur_B[node]) + 1 + return cur_mes - z = -nb_mes + 2 * idx - 2 - g = B.T.dot(z) / nb_mes +def pre_process_trees(tree_list, length_list, measures): + nx = get_backend(length_list, measures) - cur_mes -= step * g + 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_to_array([B.dot(mes[i]) for i in range(nb_mes)]) + + 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 sliced_fixed_support_tree_barycenter( + tree_list, length_list, measures, nb_itr=100, step=0.01, tol=1e-5 +): + """ + Parameters + ----------- + tree_list : array_like, shape (t, n) + length_list : array_like, shape (t, n) + measures : array_like, shape (t, m, k) + """ + + nx = get_backend(length_list, measures) + + nb_leafs = measures.shape[2] + + cur_mes = nx.ones(nb_leafs) / nb_leafs + + prepared_trees = pre_process_trees(tree_list, length_list, measures) + + nb_tree = len(prepared_trees) + + for itr in range(nb_itr): + old_mes = cur_mes.copy() + 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.dot(cur_mes) + + g_tree = get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes) + g_total += g_tree + + g_mean = g_total / nb_tree + + g_mean /= np.linalg.norm(g_mean, ord=2) + + cur_mes -= step * g_mean cur_mes = proj_simplex(cur_mes) + if np.linalg.norm(cur_mes - old_mes) < tol: + break + return cur_mes From 515ed8a49a0f89b97a25e856bf5cc6de61c3f2f0 Mon Sep 17 00:00:00 2001 From: thibaut-germain <130152047+thibaut-germain@users.noreply.github.com> Date: Tue, 7 Jul 2026 08:40:46 +0200 Subject: [PATCH 10/13] [MRG] Proximal point solver for batched optimal transport (#832) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * correction SGOT cost matrix, added to contributor list, move sgot example to other * updated graphs * fix plots * update example * fix releases.md * reformating * updating format * updated format * updated format * added PR and references * aligned usage of solve and solve_batch + added batch tests * fixed example * Apply suggestions from code review Co-authored-by: Rémi Flamary --------- Co-authored-by: Sienna O'Shea Co-authored-by: Rémi Flamary --- README.md | 6 +- RELEASES.md | 3 +- examples/backends/plot_ot_batch.py | 42 +++++-- ot/batch/__init__.py | 3 + ot/batch/_linear.py | 174 ++++++++++++++++++++++------- ot/batch/_utils.py | 154 +++++++++++++++++++++++++ test/batch/test_solve_batch.py | 170 +++++++++++++++------------- 7 files changed, 423 insertions(+), 129 deletions(-) diff --git a/README.md b/README.md index 950a509e9..87e5b9732 100644 --- a/README.md +++ b/README.md @@ -471,4 +471,8 @@ Artificial Intelligence. \[90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. -\[91] Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., & Courty, N. (2021). [Minibatch optimal transport distances; analysis and applications](https://arxiv.org/pdf/2101.01792). arXiv preprint arXiv:2101.01792. \ No newline at end of file +\[91] Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., & Courty, N. (2021). [Minibatch optimal transport distances; analysis and applications](https://arxiv.org/pdf/2101.01792). arXiv preprint arXiv:2101.01792. + +\[92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). +[A fast proximal point method for computing exact wasserstein distance.](https://proceedings.mlr.press/v115/xie20b/xie20b.pdf) In Uncertainty in artificial intelligence (pp. 433-453). PMLR. + diff --git a/RELEASES.md b/RELEASES.md index dad6fa306..5fb40e863 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -30,7 +30,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Add a numerically stable log-domain solver for entropic partial Wasserstein, selectable via the new `method` parameter of `entropic_partial_wasserstein` (`method='sinkhorn_log'`) or directly through `entropic_partial_wasserstein_logscale` (Issue #723) - Add cost functions between linear operators following [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920), - implemented in `ot.sgot` (PR #792) + implemented in `ot.sgot` (PR #792, PR #830) - Add batch FUGW loss to `ot.batch` and fix issues in some default parameters in the batch module (PR #775) - Wrapper for barycenter solvers with free support `ot.solvers.bary_free_support` (PR #730) - Build wheels on ubuntu ARM to avoid QEMU emulation (PR #818) @@ -38,6 +38,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Update the geomloss wrapper to the new version and API (PR #826) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Reorganize all tests per backend (PR #828) +- Implemented batch proximal point solver for OT problems `ot.batch.proximal_bregman_log_plan_batch` function and updated wrapper functions `ot.batch.solve_batch` and `ot.batch.solve_sample_batch` (PR #832) - Implement debiased OT solvers in `ot.solve_sample`. diff --git a/examples/backends/plot_ot_batch.py b/examples/backends/plot_ot_batch.py index dea2657e1..54573c5a4 100644 --- a/examples/backends/plot_ot_batch.py +++ b/examples/backends/plot_ot_batch.py @@ -13,6 +13,7 @@ """ # Author: Paul Krzakala +# Thibaut Germain # License: MIT License # sphinx_gallery_thumbnail_number = 1 @@ -75,25 +76,52 @@ # This is simple but inefficient for large batches. # # Instead, you can use :func:`ot.batch.solve_batch`, which solves all -# problems in parallel. +# problems in parallel. Several methods are available: ["sinkhorn", "log_sinkhorn"] +# which solve regularized OT problems, and ["proximal"] which +# solves regularized and unregularized OT problem using a proximal point scheme. +# By default, the method is set to "auto" which automatically selects the appropriate +# method based on the value of `reg`. If `reg` is None or 0, the proximal point method +# is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. + +max_iter = 10000 +tol = 1e-4 + +# Classical OT problem +## Naive approach +results_values_list = [] +for i in range(n_problems): + res = ot.solve(M_list[i], max_iter=max_iter, tol=tol) + results_values_list.append(res.value_linear) -reg = 1.0 -max_iter = 100 -tol = 1e-3 +## Batched approach +results_batch = ot.solve_batch(M=M_batch, max_iter=max_iter, tol=tol) +results_values_batch = results_batch.value_linear -# Naive approach +exact_validated = np.allclose( + np.array(results_values_list), results_values_batch, atol=tol * 10 +) + +# Entropic regularized OT problem +## Naive approach +reg = 1.0 results_values_list = [] for i in range(n_problems): res = ot.solve(M_list[i], reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy") results_values_list.append(res.value_linear) -# Batched approach +## Batched approach results_batch = ot.solve_batch( M=M_batch, reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy" ) results_values_batch = results_batch.value_linear -assert np.allclose(np.array(results_values_list), results_values_batch, atol=tol * 10) +entropic_validated = np.allclose( + np.array(results_values_list), results_values_batch, atol=tol * 10 +) + +print( + f"Exact solve vs proximal batch close: {exact_validated} \nSinkhorn solve vs Sinkhorn solve_batch close: {entropic_validated}" +) ############################################################################# # diff --git a/ot/batch/__init__.py b/ot/batch/__init__.py index d21019b98..93b998cff 100644 --- a/ot/batch/__init__.py +++ b/ot/batch/__init__.py @@ -5,6 +5,7 @@ # Author: Remi Flamary # Paul Krzakala +# Thibaut Germain # # License: MIT License @@ -25,6 +26,7 @@ bregman_log_projection_batch, bregman_projection_batch, entropy_batch, + proximal_bregman_log_plan_batch, ) __all__ = [ @@ -40,4 +42,5 @@ "loss_quadratic_batch", "loss_quadratic_samples_batch", "tensor_batch", + "proximal_bregman_log_plan_batch", ] diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index 1a9ec1955..3c07ae639 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -5,6 +5,7 @@ # Author: Remi Flamary # Paul Krzakala +# Thibaut Germain # # License: MIT License @@ -14,9 +15,17 @@ bregman_log_projection_batch, bregman_projection_batch, entropy_batch, + proximal_bregman_log_plan_batch, ) +solve_batch_method_lst = ["auto", "proximal", "log_sinkhorn", "sinkhorn"] + +solve_batch_reg_type_lst = ["kl", "entropy"] + +solve_batch_grad_lst = ["detach", "autodiff", "last_step", "envelope"] + + def dist_lp_batch(X, Y, p=2, q=1, nx=None): r"""Computes the cost matrix for a batch of samples using the Lp norm. @@ -235,23 +244,38 @@ def dist_batch( def solve_batch( M, - reg, + reg=None, a=None, b=None, max_iter=1000, tol=1e-5, - solver="log_sinkhorn", + method="auto", + inner_iter=1, + inner_reg=1e-3, reg_type="entropy", grad="envelope", ): - r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. + r""" + Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object. + + The function solves in parallel a batch of optimal transport problems: + + .. math:: + \begin{aligned} + \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{M} \rangle_F + \textit{reg} \cdot R(\mathbf{T}) \\ + \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ + & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ + & \mathbf{T} \geq 0 + \end{aligned} + + The problem is solved with either a proximal point method :ref:`[92] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function automatically selects the appropriate method based on the value of `reg`. if `reg` is None or 0, the proximal point method is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. Parameters ---------- M : array-like, shape (B, ns, nt) Cost matrix reg : float - Regularization parameter for entropic regularization + Regularization parameter. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -260,18 +284,22 @@ def solve_batch( Maximum number of iterations tol : float Tolerance for convergence - solver: str - Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. + method: str + Method to use, either 'auto', 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'auto'. + inner_iter : int + Number of inner Bregman iterations for the proximal method. Default is 1. + inner_reg : float + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-3. reg_type : str, optional Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional - Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for - Sinkhorn solver. By default 'autodiff' provides gradients wrt all - outputs (`plan, value, value_linear`) but with important memory cost. - 'envelope' provides gradients only for `value` and and other outputs are - detached. This is useful for memory saving when only the value is needed. 'last_step' provides - gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. - 'detach' does not compute the gradients for the Sinkhorn solver. + Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. + 'detach' does not compute the gradients. + 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last method iteration, useful for memory saving. + 'envelope' provides gradients only for `value`. + Default is 'envelope'. + Returns ------- @@ -292,19 +320,57 @@ def solve_batch( >>> X = np.random.randn(5, 10, 3) # 5 batches of 10 samples in 3D >>> Y = np.random.randn(5, 15, 3) # 5 batches of 15 samples in 3D >>> M = dist_batch(X, Y, metric="euclidean") # Compute cost matrices + >>> p_result = solve_batch(M) # Uses proximal method >>> reg = 0.1 - >>> result = solve_batch(M, reg) - >>> result.plan.shape # Optimal transport plans for each batch + >>> s_result = solve_batch(M, reg, method="log_sinkhorn") # Uses Sinkhorn method + >>> s_result.plan.shape # Optimal transport plans for each batch (5, 10, 15) - >>> result.value.shape # Optimal transport values for each batch + >>> s_result.value.shape # Optimal transport values for each batch (5,) See Also -------- ot.batch.dist_batch : batched cost matrix computation for computing M. - ot.solve : non-batched version of the OT solver. + ot.solve : non-batched version of the solve_batch function. + + .. _references-batch-solver: + Reference + ---------- + .. [92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + A fast proximal point method for computing exact wasserstein distance. + In Uncertainty in artificial intelligence (pp. 433-453). PMLR. + + .. [2] M. Cuturi, Sinkhorn Distances : Lightspeed Computation + of Optimal Transport, Advances in Neural Information Processing + Systems (NIPS) 26, 2013 """ + if method not in solve_batch_method_lst: + raise ValueError( + f"Unknown method: {method}. Must be one of {solve_batch_method_lst}." + ) + + if reg_type not in solve_batch_reg_type_lst: + raise ValueError( + f"Unknown reg_type: {reg_type}. Must be one of {solve_batch_reg_type_lst}." + ) + + if grad not in solve_batch_grad_lst: + raise ValueError( + f"Unknown grad: {grad}. Must be one of {solve_batch_grad_lst}." + ) + + if method in ["sinkhorn", "log_sinkhorn"] and (reg is None or reg <= 0): + raise ValueError( + "Sinkhorn methods require a strictly positive reg parameter. Please provide a valid reg value." + ) + + if method == "auto": + if reg is None or reg == 0: + method = "proximal" + else: + method = "log_sinkhorn" + nx = get_backend(a, b, M) B, n, m = M.shape @@ -314,18 +380,29 @@ def solve_batch( if b is None: b = nx.ones((B, m), type_as=M) / m - if solver == "log_sinkhorn": + if method == "log_sinkhorn": K = -M / reg out = bregman_log_projection_batch( K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad ) - elif solver == "sinkhorn": + if method == "sinkhorn": K = nx.exp(-M / reg) out = bregman_projection_batch( K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad ) - else: - raise ValueError(f"Unknown solver: {solver}") + if method == "proximal": + out = proximal_bregman_log_plan_batch( + M, + a, + b, + nx=nx, + reg=reg, + inner_reg=inner_reg, + max_iter=max_iter, + tol=tol, + inner_iter=inner_iter, + grad=grad, + ) T = out["T"] @@ -336,13 +413,15 @@ def solve_batch( T = nx.detach(T) value_linear = loss_linear_batch(M, T) - if reg_type.lower() == "entropy": + if reg_type == "entropy" and reg is not None: entr = -entropy_batch(T, nx=nx) value = value_linear + reg * entr - elif reg_type.lower() == "kl": + elif reg_type == "kl" and reg is not None: ref = nx.einsum("bi,bj->bij", a, b) kl = nx.sum(T * nx.log(T / ref + 1e-16), axis=(1, 2)) value = value_linear + reg * kl + else: + value = value_linear log = {"n_iter": out["n_iters"]} res = OTResult( @@ -360,29 +439,36 @@ def solve_batch( def solve_sample_batch( X_a, X_b, - reg, + reg=None, a=None, b=None, metric="sqeuclidean", p=2, max_iter=1000, tol=1e-5, - solver="log_sinkhorn", + method="auto", + inner_iter=1, + inner_reg=1e-3, reg_type="entropy", grad="envelope", ): - r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. + r""" + Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object computed from batches of source and target samples. + + The problem is solved with either a proximal point method :ref:`[91] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function automatically selects the appropriate method based on the value of `reg`. if `reg` is None or 0, the proximal point method is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. Parameters ---------- - M : array-like, shape (B, ns, nt) - Cost matrix - reg : float - Regularization parameter for entropic regularization + X_a : array-like, shape (B, ns, d) + Samples from source distribution + X_b : array-like, shape (B, nt, d) + Samples from target distribution metric : str, optional 'sqeuclidean', 'euclidean', 'minkowski' or 'kl' p : float, optional p-norm for the Minkowski metrics. Default value is 2. + reg : float + Regularization parameter. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -391,18 +477,21 @@ def solve_sample_batch( Maximum number of iterations tol : float Tolerance for convergence - solver: str - Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. + method: str + Method to use, either 'auto', 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'auto'. + inner_iter : int + Number of inner Bregman iterations for the proximal method. Default is 1. + inner_reg : float + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-3. reg_type : str, optional Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional - Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for - Sinkhorn solver. By default 'autodiff' provides gradients wrt all - outputs (`plan, value, value_linear`) but with important memory cost. - 'envelope' provides gradients only for `value` and and other outputs are - detached. This is useful for memory saving when only the value is needed. 'last_step' provides - gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. - 'detach' does not compute the gradients for the Sinkhorn solver. + Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. + 'detach' does not compute the gradients. + 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last method iteration, useful for memory saving. + 'envelope' provides gradients only for `value`. + Default is 'envelope'. Returns ------- @@ -418,10 +507,11 @@ def solve_sample_batch( See Also -------- - ot.batch.solve_batch : solver for computing the optimal T from arbitrary cost matrix M. + ot.batch.solve_batch : function for computing the optimal T from arbitrary cost matrix M. """ M = dist_batch(X_a, X_b, metric=metric, p=p) + return solve_batch( M, reg, @@ -429,7 +519,9 @@ def solve_sample_batch( b=b, max_iter=max_iter, tol=tol, - solver=solver, + method=method, + inner_iter=inner_iter, + inner_reg=inner_reg, reg_type=reg_type, grad=grad, ) diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index ce5db4664..8bc96c3b9 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -1,3 +1,14 @@ +# -*- coding: utf-8 -*- +""" +Utility functions for batch operations in optimal transport. +""" + +# Author: Remi Flamary +# Paul Krzakala +# Thibaut Germain +# +# License: MIT License + from ot.backend import get_backend @@ -288,3 +299,146 @@ def bregman_log_projection_batch( log_potentials = (u, v) return {"T": T, "log_T": log_T, "n_iters": n_iters, "potentials": log_potentials} + + +def proximal_bregman_log_plan_batch( + C, + a=None, + b=None, + nx=None, + reg=None, + inner_reg=1e-3, + max_iter=10000, + tol=1e-5, + inner_iter=1, + grad="detach", +): + r""" + Returns optimal transport plans for a batch of cost matrices :math:`\mathbf{C}` using a Bregman divergence based proximal point method :ref:`[92] `. + + This function solves the following optimization problem: + + .. math:: + \begin{aligned} + \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{C} \rangle_F + \textit{reg} \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} \\ + \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ + & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ + & \mathbf{T} \geq 0 + \end{aligned} + + The optimal transport plans are computed iteratively with a proximal point method based on a Bregman divergence where each iteration involves solving a Bregman projection problem: + + .. math:: + \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + (\textit{reg} + \textit{inner\_reg}) \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} + + Denoting :math:`\mathbf{K}^{(k)} = - (\mathbf{C} + \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)})/(\textit{reg} + \textit{inner\_reg})`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: + + .. math:: + \mathbf{u}^{(i+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K}^{(k)} + \mathbf{v}^{(i)}) + + \mathbf{v}^{(i+1)} = \log(\mathbf{b}) - \text{LSE}((\mathbf{K}^{(k)})^T + \mathbf{u}^{(i)}) + + where LSE denotes the log-sum-exp operation. + + Parameters + ---------- + C : array-like, shape (B, n, m) + Cost matrix for each problem in the batch. + a : array-like, shape (B, n), optional + Source distribution for each problem. If None, uniform distribution is used. + b : array-like, shape (B, m), optional + Target distribution for each problem. If None, uniform distribution is used. + nx : backend object, optional + Numerical backend to use for computations. If None, the default backend is used. + max_iter : int, optional + Maximum number of iterations. + inner_iter : int, optional + Number of inner iterations for updating the dual variables u and v. Default is 1. + inner_reg : float, optional + Regularization parameter for the Bregman divergence. Default is 1e-3. + tol : float, optional + Tolerance for convergence. The solver stops when the maximum change in + the dual variables is below this value. + grad : str, optional + Gradient computation mode: 'detach', 'autodiff', or 'last_step'. + + Returns + ------- + dict + Dictionary containing: + - 'T' : array-like, shape (B, n, m) + Optimal transport plan for each problem. + - 'log_T' : array-like, shape (B, n, m) + Logarithm of the optimal transport plan for each problem. + - 'potentials' : tuple of array-like, shapes ((B, n), (B, m)) + Log-scaling factors (u, v). + - 'n_iters' : int + Number of iterations performed. + + Example + -------- + >>> import numpy as np + >>> from ot.batch import proximal_bregman_log_plan_batch + >>> # Create batch of affinity matrices + >>> C = np.random.rand(5, 10, 15) # 5 problems, 10x15 cost matrices + >>> result = proximal_bregman_log_plan_batch(C) + >>> T = result['T'] # Shape (5, 10, 15) + + .. _references-OT-bregman-proximal-point: + Reference + ---------- + .. [92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + A fast proximal point method for computing exact wasserstein distance. + In Uncertainty in artificial intelligence (pp. 433-453). PMLR. + """ + + if nx is None: + nx = get_backend(a, b, C) + + B, n, m = C.shape + + if a is None: + a = nx.ones((B, n)) / n + if b is None: + b = nx.ones((B, m)) / m + + if reg is None: + reg = 0.0 + + u = nx.zeros((B, n), type_as=C) + v = nx.zeros((B, m), type_as=C) + + loga = nx.log(a) + logb = nx.log(b) + + if grad == "detach": + C = nx.detach(C) + elif grad == "last_step": + C_, C = C.clone(), nx.detach(C) + + log_T = nx.zeros(C.shape, type_as=C) + for n_iters in range(max_iter): + K_proj = -(C + inner_reg * log_T) / (reg + inner_reg) + for _ in range(inner_iter): + u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) + v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) + log_T = K_proj + u[:, :, None] + v[:, None, :] + + # Check convergence once every 10 iterations + if n_iters % 10 == 0: + T = nx.exp(log_T) + marginal = nx.sum(T, axis=2) + err = nx.max(norm_batch(marginal - a)) + if err < tol: + break + + if grad == "last_step": + K_proj = -(C_ + inner_reg * log_T) / (reg + inner_reg) + for _ in range(inner_iter): + u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) + v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) + log_T = K_proj + u[:, :, None] + v[:, None, :] + T = nx.exp(log_T) + log_potentials = (u, v) + + return {"T": T, "log_T": log_T, "n_iters": n_iters, "potentials": log_potentials} diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 17d459a43..08e5ff81e 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -3,6 +3,7 @@ # Author: Remi Flamary # Paul Krzakala # Sonia Mazelet +# Thibaut Germain # # License: MIT License @@ -15,50 +16,85 @@ loss_linear_batch, bregman_projection_batch, bregman_log_projection_batch, + proximal_bregman_log_plan_batch, ) from ot import solve +from contextlib import nullcontext import pytest from ot.backend import torch -@pytest.mark.parametrize("solver", ["sinkhorn", "log_sinkhorn"]) +@pytest.mark.parametrize("reg", [None, 0, 1e-0]) +@pytest.mark.parametrize("method", ["auto", "proximal", "sinkhorn", "log_sinkhorn"]) @pytest.mark.parametrize("reg_type", ["kl", "entropy"]) -def test_solve_batch(solver, reg_type): +def test_solve_batch_vs_solve(reg, method, reg_type): """Check that solve_batch gives the same results as solve for each instance in the batch.""" - batchsize = 4 - n = 16 - rng = np.random.RandomState(0) - - M = rng.rand(batchsize, n, n) - - reg = 0.1 - max_iter = 10000 - tol = 1e-5 - res = solve_batch( - M, - a=None, - b=None, + should_fail = method in ["sinkhorn", "log_sinkhorn"] and (reg is None or reg <= 0) + + ctx = pytest.raises(Exception) if should_fail else nullcontext() + + with ctx: + tol = 1e-4 + batchsize = 3 + n = 5 + d = 7 + rng = np.random.RandomState(0) + C = rng.rand(batchsize, n, d) + + base_plan = np.zeros((batchsize, n, d)) + base_value = np.zeros(batchsize) + for i in range(batchsize): + C_i = C[i] + res_i = solve(C_i, reg=reg, tol=tol, reg_type=reg_type) + base_plan[i] = res_i.plan + base_value[i] = res_i.value_linear + + res = solve_batch( + C, + max_iter=10000, + tol=tol, + grad="detach", + reg=reg, + method=method, + reg_type=reg_type, + inner_reg=1e-3, + ) + plan = res.plan + value = res.value_linear + np.testing.assert_allclose(plan, base_plan, atol=tol * 10) + np.testing.assert_allclose(value, base_value, atol=tol * 10) + + +@pytest.mark.parametrize("reg", [None, 0, 1e-0]) +@pytest.mark.parametrize("inner_iter", [1, 5, 10]) +def test_backend_proximal_bregman_log_plan_batch(nx, reg, inner_iter): + tol = 1e-4 + batchsize = 3 + n = 5 + d = 7 + rng = np.random.RandomState(0) + C = rng.rand(batchsize, n, d) + res = proximal_bregman_log_plan_batch( + nx.from_numpy(C), reg=reg, - max_iter=max_iter, + inner_reg=1e-3, + max_iter=10000, tol=tol, - solver=solver, - reg_type=reg_type, + inner_iter=inner_iter, grad="detach", ) - plan_batch = res.plan - values_batch = res.value_linear - + plan = nx.to_numpy(res["T"]) for i in range(batchsize): - M_i = M[i] + C_i = C[i] res_i = solve( - M_i, a=None, b=None, reg=reg, max_iter=max_iter, tol=tol, reg_type=reg_type + C_i, + reg=reg, + tol=tol, ) plan_i = res_i.plan - value_i = res_i.value_linear - np.testing.assert_allclose(plan_i, plan_batch[i], atol=1e-05) - np.testing.assert_allclose(value_i, values_batch[i], atol=1e-4) + np.testing.assert_allclose(plan_i, plan[i], atol=tol * 10) def test_bregman_batch(): @@ -78,9 +114,10 @@ def test_bregman_batch(): @pytest.mark.parametrize("metric", ["sqeuclidean", "euclidean", "minkowski", "kl"]) -def test_metrics(metric): +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_sample_solve_batch_vs_solve_batch(metric, method): """Check that all functions run without error.""" - + tol = 1e-5 batchsize = 2 n = 4 d = 2 @@ -93,11 +130,10 @@ def test_metrics(metric): is_positive = M >= 0 np.testing.assert_equal(is_positive.all(), True) - # Solve batch - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - # Solve sample batch - res = solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, metric=metric) + res = solve_sample_batch( + X, X, reg=0.1, max_iter=10, tol=tol, metric=metric, method=method + ) # Compute loss loss = res.value_linear # loss given by solver @@ -105,34 +141,35 @@ def test_metrics(metric): loss3 = loss_linear_samples_batch( X, X, res.plan, metric=metric ) # recompute loss from plan and samples - np.testing.assert_allclose(loss, loss2, atol=1e-5) - np.testing.assert_allclose(loss, loss3, atol=1e-5) + np.testing.assert_allclose(loss, loss2, atol=tol * 10) + np.testing.assert_allclose(loss, loss3, atol=tol * 10) @pytest.mark.skipif(not torch, reason="torch not installed") @pytest.mark.parametrize("grad", ["detach", "envelope", "autodiff", "last_step"]) -def test_gradients_torch(grad): +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_gradients_torch(grad, method): """Check that all gradient methods run without error.""" batchsize = 2 n = 4 d = 2 - for solver in ["sinkhorn", "log_sinkhorn"]: - X = torch.randn((batchsize, n, d), requires_grad=True) - M = dist_batch(X, X) - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, solver=solver) - loss = res.value_linear.sum() - loss_plan = res.plan.sum() - if grad == "detach": - assert loss.grad == None - elif grad == "envelope": - loss.backward() - assert X.grad is not None - elif grad in ["autodiff", "last_step"]: - loss_plan.backward() - assert X.grad is not None - - -def test_backend(nx): + X = torch.randn((batchsize, n, d), requires_grad=True) + M = dist_batch(X, X) + res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, method=method) + loss = res.value_linear.sum() + loss_plan = res.plan.sum() + if grad == "detach": + assert loss.grad == None + elif grad == "envelope": + loss.backward() + assert X.grad is not None + elif grad in ["autodiff", "last_step"]: + loss_plan.backward() + assert X.grad is not None + + +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_backend(nx, method): """Check that all gradient methods run without error.""" batchsize = 2 n = 4 @@ -140,30 +177,5 @@ def test_backend(nx): X = np.random.randn(batchsize, n, d) X = nx.from_numpy(X) M = dist_batch(X, X) - solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5) - - -def test_metric_default_parameters(): - """Check that all functions with default parameters run without error.""" - - batchsize = 2 - n = 4 - d = 2 - rng = np.random.RandomState(0) - X = rng.rand(batchsize, n, d) - M = dist_batch(X, X) - is_positive = M >= 0 - np.testing.assert_equal(is_positive.all(), True) - - # Solve batch - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - - # Solve sample batch - res = solve_sample_batch(X, X, reg=0.1) - - # Compute loss - loss_linear_batch(M, res.plan) # recompute loss from plan - loss_linear_samples_batch(X, X, res.plan) # recompute loss from plan and samples - assert np.isfinite(loss_linear_batch(M, res.plan)).all() - assert np.isfinite(loss_linear_samples_batch(X, X, res.plan)).all() + solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, method=method) + solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, method=method) From d4076995599fecb6bd61a056bdc94a767bb14bc9 Mon Sep 17 00:00:00 2001 From: CodingSelim <139068169+CodingSelim@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:31:05 +0300 Subject: [PATCH 11/13] [MRG] Fix label-aware cost correction in ot.da (closes #664) (#833) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix label-aware cost correction in ot.da the cost correction that pushes apart differently-labeled samples was computed backwards. it used missing_ys = (ys == -1) and the product of the two missing masks, so the large cost was applied only where both labels are missing, and never where two labeled samples have different labels. with all labels known the correction was a no-op, so ys/yt had no effect on the transport (see #664). switched to present masks (ys != -1) so the correction applies exactly to labeled source/target pairs whose labels differ, matching the original pre-vectorized loop. restored the semisupervised tests that had been flipped to assert the buggy no-op (n_unsup == n_semisup) back to asserting the cost actually changes, and added a regression test. closes #664 * add PR number to releases entry * Apply suggestion from @rflamary * use MISSING_LABEL constant instead of bare -1 --------- Co-authored-by: Rémi Flamary --- RELEASES.md | 1 + ot/da.py | 26 +++++++++++++++----------- test/test_da.py | 43 ++++++++++++++++++++++++++++++++++++++----- 3 files changed, 54 insertions(+), 16 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 5fb40e863..5d9c5b61a 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -44,6 +44,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s #### Closed issues +- Fix label-aware cost correction in `ot.da` transport classes: the large cost was applied to unlabeled pairs instead of labeled pairs with different labels, so `ys`/`yt` had no effect when all labels were known (PR #833, Issue #664) - Mitigate NaN regime of `entropic_partial_wasserstein` at small `reg` via a new log-domain solver, reachable with `entropic_partial_wasserstein(..., method='sinkhorn_log')` (Issue #723; the default `method='sinkhorn'` path is unchanged — callers opt into the log-domain variant) - Fix NumPy 2.x compatibility in Brenier potential bounds (PR #788) - Fix MSVC Windows build by removing **restrict** keyword (PR #788) diff --git a/ot/da.py b/ot/da.py index 527999cba..7b4ed7ba2 100644 --- a/ot/da.py +++ b/ot/da.py @@ -48,6 +48,9 @@ joint_OT_mapping_kernel, ) +# value used in ys/yt to mark a sample whose label is unknown +MISSING_LABEL = -1 + def sinkhorn_lpl1_mm( a, @@ -598,23 +601,24 @@ class label if self.limit_max != np.inf: self.limit_max = self.limit_max * nx.max(self.cost_) - # missing_labels is a (ns, nt) matrix of {0, 1} such that - # the cells (i, j) has 0 iff either ys[i] or yt[j] is masked - missing_ys = (ys == -1) + nx.zeros(ys.shape, type_as=ys) - missing_yt = (yt == -1) + nx.zeros(yt.shape, type_as=yt) - missing_labels = missing_ys[:, None] @ missing_yt[None, :] - # labels_match is a (ns, nt) matrix of {True, False} such that - # the cells (i, j) has False if ys[i] != yt[i] - label_match = (ys[:, None] - yt[None, :]) != 0 + # present_labels is a (ns, nt) matrix of {0, 1} such that + # the cell (i, j) is 1 iff both ys[i] and yt[j] are labeled + # (i.e. neither is masked with MISSING_LABEL) + present_ys = (ys != MISSING_LABEL) + nx.zeros(ys.shape, type_as=ys) + present_yt = (yt != MISSING_LABEL) + nx.zeros(yt.shape, type_as=yt) + present_labels = present_ys[:, None] @ present_yt[None, :] + # label_mismatch is a (ns, nt) matrix of {True, False} such that + # the cell (i, j) is True if ys[i] != yt[j] + label_mismatch = (ys[:, None] - yt[None, :]) != 0 # cost correction is a (ns, nt) matrix of {-Inf, float, Inf} such # that he cells (i, j) has -Inf where there's no correction necessary - # by 'correction' we mean setting cost to a large value when - # labels do not match + # by 'correction' we mean setting cost to a large value when two + # labeled samples have different labels # we suppress potential RuntimeWarning caused by Inf multiplication # (as we explicitly cover potential NANs later) with warnings.catch_warnings(): warnings.simplefilter("ignore", category=RuntimeWarning) - cost_correction = label_match * missing_labels * self.limit_max + cost_correction = label_mismatch * present_labels * self.limit_max # this operation is necessary because 0 * Inf = NAN # thus is irrelevant when limit_max is finite cost_correction = nx.nan_to_num(cost_correction, -np.inf) diff --git a/test/test_da.py b/test/test_da.py index bb548e27f..df224a247 100644 --- a/test/test_da.py +++ b/test/test_da.py @@ -166,7 +166,7 @@ def test_sinkhorn_lpl1_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -258,7 +258,7 @@ def test_sinkhorn_l1l2_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -356,7 +356,7 @@ def test_sinkhorn_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -472,7 +472,7 @@ def test_unbalanced_sinkhorn_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -571,7 +571,7 @@ def test_emd_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -586,6 +586,39 @@ def test_emd_transport_class(nx): ) +@pytest.skip_backend("tf") +def test_semisupervised_cost_correction(nx): + # gh-664: the label-aware cost correction must push apart labeled source + # and target samples with *different* labels, and must leave pairs + # involving an unlabeled (-1) sample untouched. + rng = np.random.RandomState(0) + Xs = rng.randn(6, 2) + Xt = rng.randn(6, 2) + ys = np.array([0, 0, 0, 1, 1, 1]) + yt = np.array([0, 1, 0, 1, 0, 1]) + Xs, ys, Xt, yt = nx.from_numpy(Xs, ys, Xt, yt) + + # all labels known: different-label pairs get the (scaled) limit_max, + # same-label pairs keep their original, smaller cost + otda = ot.da.EMDTransport(limit_max=10) + otda.fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt) + cost = nx.to_numpy(otda.cost_) + limit = float(nx.to_numpy(otda.limit_max)) + ys_np, yt_np = nx.to_numpy(ys), nx.to_numpy(yt) + mismatch = ys_np[:, None] != yt_np[None, :] + assert np.all(cost[mismatch] >= limit - 1e-9), "different labels not penalized" + assert np.all(cost[~mismatch] < limit), "same labels wrongly penalized" + + # semi-supervised: unlabeled (-1) targets must never be forbidden + yt_semi = nx.from_numpy(np.array([0, -1, 0, -1, 0, -1])) + otda2 = ot.da.EMDTransport(limit_max=10) + otda2.fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt_semi) + cost2 = nx.to_numpy(otda2.cost_) + limit2 = float(nx.to_numpy(otda2.limit_max)) + unlabeled = nx.to_numpy(yt_semi) == -1 + assert np.all(cost2[:, unlabeled] < limit2), "unlabeled targets wrongly forbidden" + + @pytest.skip_backend("jax") @pytest.skip_backend("tf") @pytest.mark.parametrize("kernel", ["linear", "gaussian"]) From 3f6183bf8b4c01805a5885474137f15110b7e64f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Tue, 7 Jul 2026 12:41:35 +0200 Subject: [PATCH 12/13] [MRG] Release 0.9.7 (#829) * nx version + update actions * add run on arm * use slim ubuntu * back to linux latest * reamp all tests * fix yaml * fix yaml * fix yaml * try it * validate yaml * speeup backen tests * remove doctest for beckend tests * remove cleanup space * separate doctests * fix cdoctest with proper cnftest * fix helpers * fix linux-torch test * remove torch version * fix doc build * fix doc build * fix doctests * move stuff around properly * gix doctest for the last time? * stuff * change version * pdate reelase + readme * update cff + readme * big rewrite release --- CITATION.cff | 1 + README.md | 18 +++++++++++------- RELEASES.md | 41 ++++++++++++++++++++++++++++++++++++----- ot/__init__.py | 2 +- 4 files changed, 49 insertions(+), 13 deletions(-) diff --git a/CITATION.cff b/CITATION.cff index 513398eb7..d7e4c56ab 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -109,3 +109,4 @@ keywords: - gromov-wasserstein license: MIT version: 0.9.7 +date-released: 2026-07-10 diff --git a/README.md b/README.md index 87e5b9732..70092d9ad 100644 --- a/README.md +++ b/README.md @@ -96,7 +96,7 @@ using the following references from the current version and from our [JMLR paper](https://jmlr.org/papers/v22/20-451.html): ``` -Flamary R., Vincent-Cuaz C., Courty N., Gramfort A., Kachaiev O., Quang Tran H., David L., Bonet C., Cassereau N., Gnassounou T., Tanguy E., Delon J., Collas A., Mazelet S., Chapel L., Kerdoncuff T., Yu X., Feickert M., Krzakala P., Liu T., Fernandes Montesuma E. POT Python Optimal Transport (version 0.9.5). URL: https://github.com/PythonOT/POT +Flamary R., Vincent-Cuaz C., Courty N., Gramfort A., Kachaiev O., Quang Tran H., David L., Bonet C., Cassereau N., Gnassounou T., Tanguy E., Delon J., Collas A., Mazelet S., Chapel L., Kerdoncuff T., Yu X., Feickert M., Krzakala P., Liu T., Fernandes Montesuma E., Neike N., Genest B., Coeurjolly D., Germain T., O'Shea S., Corneli M., Genans F. (2026). POT Python Optimal Transport (version 0.9.7). DOI: 10.5281/zenodo.17161062 URL: https://github.com/PythonOT/POT Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, Titouan Vayer, POT Python Optimal Transport library, Journal of Machine Learning Research, 22(78):1−8, 2021. URL: https://pythonot.github.io/ ``` @@ -104,13 +104,17 @@ Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Bo In Bibtex format: ```bibtex -@misc{flamary2024pot, - author = {Flamary, R{\'e}mi and Vincent-Cuaz, C{\'e}dric and Courty, Nicolas and Gramfort, Alexandre and Kachaiev, Oleksii and Quang Tran, Huy and David, Laurène and Bonet, Cl{\'e}ment and Cassereau, Nathan and Gnassounou, Th{\'e}o and Tanguy, Eloi and Delon, Julie and Collas, Antoine and Mazelet, Sonia and Chapel, Laetitia and Kerdoncuff, Tanguy and Yu, Xizheng and Feickert, Matthew and Krzakala, Paul and Liu, Tianlin and Fernandes Montesuma, Eduardo}, - title = {POT Python Optimal Transport (version 0.9.5)}, - url = {https://github.com/PythonOT/POT}, - year = {2024} +@software{flamary2026pot, +author = {Flamary, Rémi and Vincent-Cuaz, Cédric and Courty, Nicolas and Gramfort, Alexandre and Kachaiev, Oleksii and Quang Tran, Huy and David, Laurène and Bonet, Clément and Cassereau, Nathan and Gnassounou, Théo and Tanguy, Eloi and Delon, Julie and Collas, Antoine and Mazelet, Sonia and Chapel, Laetitia and Kerdoncuff, Tanguy and Yu, Xizheng and Feickert, Matthew and Krzakala, Paul and Liu, Tianlin and Fernandes Montesuma, Eduardo and Neike, Nathan and Genest, Baptiste and Coeurjolly, David and Germain, Thibaut and O'Shea, Sienna and Corneli, Marco and Genans, Ferdinand}, +doi = {10.5281/zenodo.17161062}, +month = {7}, +title = {POT Python Optimal Transport}, +version = {0.9.7}, +url = {https://github.com/PythonOT/POT}, +year = {2026} } + @article{flamary2021pot, author = {R{\'e}mi Flamary and Nicolas Courty and Alexandre Gramfort and Mokhtar Z. Alaya and Aur{\'e}lie Boisbunon and Stanislas Chambon and Laetitia Chapel and Adrien Corenflos and Kilian Fatras and Nemo Fournier and L{\'e}o Gautheron and Nathalie T.H. Gayraud and Hicham Janati and Alain Rakotomamonjy and Ievgen Redko and Antoine Rolet and Antony Schutz and Vivien Seguy and Danica J. Sutherland and Romain Tavenard and Alexander Tong and Titouan Vayer}, title = {POT: Python Optimal Transport}, @@ -467,7 +471,7 @@ Artificial Intelligence. \[88] Bouveyron, C. & Corneli, M. (2026). [Scaling optimal transport to high-dimensional Gaussian distributions with application to domain adaptation](https://hal.science/hal-04930868v4/file/Article-OT-HDGauss-v4.pdf). Statistics and Computing 36.2 (2026): 88. -\[89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis]. Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622. +\[89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis](https://www.cs.columbia.edu/~blei/seminar/2020-representation/readings/TippingBishop1999.pdf). Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622. \[90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. diff --git a/RELEASES.md b/RELEASES.md index 5d9c5b61a..3e6014315 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -2,11 +2,43 @@ ## 0.9.7 -This new release adds support for sparse cost matrices and a new lazy exact OT solver that re-computes distances on-the-fly from coordinates, reducing memory usage from O(n×m) to O(n+m). Both implementations are backend-agnostic and preserve gradient computation for automatic differentiation. +This release contains several bug fixes and performance improvements, as well as updates to the documentation and examples and tests. We provide below a summary of the main new features and closed issues. + + +**New exact OT solvers.** This new release adds many new variants and updates for the exact OT solver. First a new lazy exact OT solver that re-computes distances on-the-fly from coordinates, reducing memory usage from O(n×m) to O(n+m) ( available in [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample) with `lazy=True)`). Another major feature is the addition of a warmstart feature for the EMD solver, allowing users to provide initial potentials to speed up convergence (with `init_potentials` in [`ot.solve`](https://pythonot.github.io/all.html#ot.solve) and [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample)). Finally the release also include a sparse solver and [`ot.solve`](https://pythonot.github.io/all.html#ot.solve) now accept sparse cost matrices and provides an OT plan whose support is included in the support of the cost matrix. All implementations are backend-agnostic and preserve gradient computation for automatic differentiation (but are solved on CPU). + +The computational time for different solvers (on a laptop CPU) are shown below: + +| n| 100| 500| 1000| 5000| +|--------------------------|:------------------------:|:------------------------:|:------------------------:|:------------------------:| +| Solve | 7.259e-04| 1.501e-02| 1.055e-01| 2.343e+00| +| Solve Warm Start | 6.970e-04| 5.413e-03| 2.291e-02| 6.036e-01| +| Solve Lazy | 1.280e-03| 3.098e-02| 1.421e-01| 3.377e+00| +| Solve Sparse 10% | 3.453e-04| 6.716e-04| 1.535e-03| 3.534e-02| + + +**BSPOT solver.** A new solver relying on [Binary Space Partitioning (BSP)](https://pythonot.github.io/master/auto_examples/plot_bsp_ot.html) has been added to compute sparse transport plans between discrete measures in loglinear time which allows for very large problems to be solved. The release also includes a new [`ot.unbalanced.uot_1d`](https://pythonot.github.io/master/gen_modules/ot.unbalanced.html#ot.unbalanced.uot_1d) solver with a Frank-Wolfe solver for unbalanced optimal transport in 1D. + +**Sliced OT pans** Finally we now have a sliced OT plan solver that can be used to compute sliced transport plans ([min-pivot sliced](https://pythonot.github.io/master/gen_modules/ot.sliced.html#ot.sliced.min_sliced_transport_plan) and [expected sliced](https://pythonot.github.io/master/gen_modules/ot.sliced.html#ot.sliced.expected_sliced_plan)) between two measures. + +**OT between dynamical systems.** A novel [Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://pythonot.github.io/master/auto_examples/others/plot_sgot.html) has been implemented in `ot.sgot`. + +**Unified API for barycenter solvers in `ot.solve_bary_sample`.** A new free support solver for barycenter solvers has been added in [`ot.solve_bary_sample`](https://pythonot.github.io/master/all.html#ot.solve_bary_sample). You can see [examples of how to use it here](https://pythonot.github.io/master/auto_examples/barycenters/plot_solve_barycenter_variants.html). + +**OT between high dimensional Gaussian distributions.** New methods to compute the linear transport map and the related 2-Wasserstein distance between high-dimensional (HD) Gaussian distributions have been added in [`ot.gaussian.bures_wasserstein_mapping_hd`](https://pythonot.github.io/master/gen_modules/ot.gaussian.html#id56) and [`ot.gaussian.bures_wasserstein_distance_hd`](https://pythonot.github.io/master/gen_modules/ot.gaussian.html#ot.gaussian.bures_wasserstein_distance_hd), respectively with empirical versions that estimate the distance from sample. + +**Batched solver for exact OT.** The batch implementations have also been improved and you can now solve exact OT problems in parallel on CPU or GPU using the new proximal point solver in functions [`ot.solve_batch`](https://pythonot.github.io/master/gen_modules/ot.batch.html#ot.batch.solve_batch) and [`ot.solve_sample_batch`](https://pythonot.github.io/master/gen_modules/ot.batch.html#ot.batch.solve_sample_batch) when `reg=0` or no provided. A new batch [loss for Fused unbalanced Gromov-Wasserstein](https://pythonot.github.io/master/gen_modules/ot.batch.html#ot.batch.loss_quadratic_batch) is also now available in `ot.batch`. + + +**New methods in unified API [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample).** The unified API function +[`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample) has also been updated to allows solving of specific problems such as BSP-OT, distance between high-dimensional (HD) Gaussian distributions and sliced and max-sliced distances. + + +**Data normalization for sliced and [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample) solvers.** The release also includes tools for data normalization and scaling which can not be used in sliced Wasserstein distance computations. A simple normalization class [`ot.utils.DataScaler`](https://pythonot.github.io/master/gen_modules/ot.utils.html#id46) has been added and supports all backends for `'standard'`, `'minmax'`, and `'l2'` methods. Finally an optional `scaler` parameter has been added to [`ot.sliced_wasserstein_distance`](https://pythonot.github.io/master/all.html#ot.sliced_wasserstein_distance), [`ot.max_sliced_wasserstein_distance`](https://pythonot.github.io/master/all.html#ot.max_sliced_wasserstein_distance) and [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample). + #### New features -- Fix reference number error introduced in PR #767 (PR #819) - Refactor lazy EMD network simplex storage to avoid dense per-arc cost, endpoint, flow, and state storage where possible, and return sparse lazy transport plans instead of materializing dense plans internally (PR #813) @@ -28,9 +60,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Add optional `scaler` parameter to `sliced_wasserstein_distance` and `max_sliced_wasserstein_distance` (PR #808) - Add SGD based semi-discrete OT solver in `ot.semidiscrete` and a gallery example. (PR #812) - Add a numerically stable log-domain solver for entropic partial Wasserstein, selectable via the new `method` parameter of `entropic_partial_wasserstein` (`method='sinkhorn_log'`) or directly through `entropic_partial_wasserstein_logscale` (Issue #723) -- Add cost functions between linear operators following - [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920), - implemented in `ot.sgot` (PR #792, PR #830) +- Add cost functions between linear operators following [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920), implemented in `ot.sgot` (PR #792, PR #830) - Add batch FUGW loss to `ot.batch` and fix issues in some default parameters in the batch module (PR #775) - Wrapper for barycenter solvers with free support `ot.solvers.bary_free_support` (PR #730) - Build wheels on ubuntu ARM to avoid QEMU emulation (PR #818) @@ -62,6 +92,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Fix entropic regularization in `gcg`(PR #817, Issue #758) - Fix documentation build on master with submodules (PR #818) - Fix failing test for unbalanced solver with generic regularization (PR #824) +- Fix reference number error introduced in PR #767 (PR #819) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Update sgot cost function and example (PR #830) diff --git a/ot/__init__.py b/ot/__init__.py index 59cd4f2d6..d867daae4 100644 --- a/ot/__init__.py +++ b/ot/__init__.py @@ -94,7 +94,7 @@ from .utils import dist, unif, tic, toc, toq -__version__ = "0.9.7.dev0" +__version__ = "0.9.7" __all__ = [ "emd", From 3dbda889ff3dadbd7cd21e8a0f5d0eabb4e1fa2f Mon Sep 17 00:00:00 2001 From: Boudjema Ali Date: Wed, 8 Jul 2026 11:53:49 +0200 Subject: [PATCH 13/13] doc --- ot/lp/__init__.py | 8 +- ot/lp/solver_tree.py | 6 +- ot/lp/tree_barycenter.py | 189 ++++++++++++--------------------------- 3 files changed, 63 insertions(+), 140 deletions(-) diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index 3aa010a1f..32bbbf550 100644 --- a/ot/lp/__init__.py +++ b/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", ] diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py index 498d4778c..079227c83 100644 --- a/ot/lp/solver_tree.py +++ b/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) diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py index f8b58c19f..6a7811580 100644 --- a/ot/lp/tree_barycenter.py +++ b/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)