[MRG] OT barycenters for generic transport costs #715

eloitanguy · 2025-01-20T16:16:09Z

Types of changes

contribute a generalisation of free_support that accepts any cost function (implements this paper)
contribute a specific version of this function in ot.gmm for fast computation of GMM barycenters
example for new barycenter function
example of GMM barycenters
add summary of features and example in README.md
added a solver for the "true" fixed-point algo which increases support sizes across iterations

For the (theoretical) fixed-point method, use method='true_fixed_point' in ot.lp.free_support_barycenter_generic_costs and for the barycentric heuristic, use method='L2_barycentric_proj'. The latter is the default, given the computational advantages and the desirable property of keeping a fixed support size.

Motivation and context / Related issue

How has this been tested (if it applies)

full coverage in test/test_ot.py and test/test_gmm.py

PR checklist

I have read the CONTRIBUTING document.
The documentation is up-to-date with the changes I made (check build artifacts).
All tests passed, and additional code has been covered with new tests.
I have added the PR and Issue fix to the RELEASES.md file.

….cvx

codecov · 2025-01-20T16:48:38Z

Codecov Report

❌ Patch coverage is 99.24812% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.14%. Comparing base (a78b197) to head (56c1499).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #715      +/-   ##
==========================================
+ Coverage   97.10%   97.14%   +0.04%     
==========================================
  Files         100      100              
  Lines       20453    20849     +396     
==========================================
+ Hits        19861    20254     +393     
- Misses        592      595       +3

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

rflamary

thanks @eloitanguy this is very nice work as usual. I have a couple comments below

rflamary · 2025-03-14T12:39:45Z

ot/lp/_barycenter_solvers.py

+        List of K arrays of measure weights, each of shape (m_k).
+    X_init : array-like
+        Array of shape (n, d) representing initial barycenter points.
+    cost_list : list of callable


could you accept a callable (if same cost for eveyone) and a list?

ot/lp/_barycenter_solvers.py

examples/barycenters/plot_free_support_barycenter_generic_cost.py

examples/barycenters/plot_gmm_barycenter.py

eloitanguy · 2025-03-18T13:01:31Z

After discussion with @rflamary we decided to expand the PR with additional features, namely two barycenter solvers:

a solver for the iterates of H of 76 with an optional ground_bary function computing barycenters between the points on different spaces, and with a user-defined notion of cost on each space $\mathcal{Y}_k$ with respect to which the barycentric projection is performed (defaults to L2, which is fast but often not accurate/meaningful depending on the spaces and costs)
another BCD solver which alternates between the plans and the the positions X (optimisation w.r.t. the positions with a pytorch optimiser)

…and examples

…gnment

eloitanguy · 2025-06-06T15:25:15Z

After some updates to the paper behind these algorithms (the paper update is soon to come), I implemented another barycenter solver which corresponds exactly to the method studied theoretically in the paper (iterates of G). I updated this PR with the additional algorithm (method='true_fixed_point'), with a comparison of the two methods in the example.

On my end, this contribution is ready for review :D

cedricvincentcuaz

Thanks a lot for this very nice PR @eloitanguy ! :)
Follows some comments on the generic barycentre solver for now. I'll revise the GMM parts next.

cedricvincentcuaz · 2025-08-07T14:03:49Z

ot/lp/_barycenter_solvers.py

+        Barycentre method: 'L2_barycentric_proj' (default) for Euclidean
+        barycentric projection, or 'true_fixed_point' for iterates using the
+        North West Corner multi-marginal gluing method.
+    stopThr : float, optional


Specify the criterion

cedricvincentcuaz · 2025-08-07T14:04:47Z

ot/lp/_barycenter_solvers.py

+    ground_bary_numItermax : int, optional
+        Maximum number of iterations for the ground barycenter solver (if auto
+        is used).
+    ground_bary_stopThr : float, optional


specify the criterion.

cedricvincentcuaz · 2025-08-07T14:07:15Z

ot/lp/_barycenter_solvers.py

+            emd(a, measure_weights[k], cost_list[k](X, measure_locations[k]))
+            for k in range(K)
+        ]
+        Y_perm = []


to remove for consistency

cedricvincentcuaz · 2025-08-07T14:08:23Z

ot/lp/_barycenter_solvers.py

+        if method == "L2_barycentric_proj":
+            a_next = a  # barycentre weights are fixed
+            for k in range(K):  # L2 barycentric projection of pi_k
+                Y_perm.append((1 / a[:, None]) * pi_list[k] @ measure_locations[k])


Y_perm = [... for k in range(K)]

cedricvincentcuaz · 2025-08-07T14:14:54Z

ot/lp/_barycenter_solvers.py

+            raise ImportError("PyTorch is required to use ground_bary=None")
+
+    X_list = [X_init] if log else []  # store the iterations
+    a_list = [nx.copy(a)] if log and method == "true_fixed_point" else []


copy needed ?

cedricvincentcuaz · 2025-08-07T16:33:40Z

ot/lp/_barycenter_solvers.py

+        m_k}`. If cost_list is a single callable, the same cost is used K times.
+    ground_bary : callable or None, optional
+        Function List(array(n, d_k)) -> array(n, d) accepting a list of K arrays
+        of shape (n\times d_K), computing the ground barycenters (broadcasted


cedricvincentcuaz · 2025-08-07T16:37:38Z

examples/barycenters/plot_free_support_barycenter_generic_cost.py

+(non-linear) projection onto a circle k, and :math:`(\lambda_k)` are weights. A
+barycenter is defined ([76]) as a minimiser of the energy :math:`V(\mu) = \sum_k
+\mathcal{T}_{c_k}(\mu, \nu_k)` where :math:`\mu` is a candidate barycenter
+measure, the measures  :math:`\nu_k` are the target measures and


rendering issue ?

cedricvincentcuaz · 2025-08-07T16:38:39Z

examples/barycenters/plot_free_support_barycenter_generic_cost.py

+
+torch.manual_seed(42)
+
+n = 136  # number of points of the of the barycentre


remove of the

cedricvincentcuaz · 2025-08-07T16:39:55Z

examples/barycenters/plot_free_support_barycenter_generic_cost.py

+n = 136  # number of points of the of the barycentre
+d = 2  # dimensions of the original measure
+K = 4  # number of measures to barycentre
+m = 50  # number of points of the measures


can we take different number of points ?

cedricvincentcuaz · 2025-08-07T16:43:56Z

examples/barycenters/plot_free_support_barycenter_generic_cost.py

+# -*- coding: utf-8 -*-
+"""
+=====================================
+OT Barycenter with Generic Costs Demo


very nice illustrations :)

eloitanguy and others added 14 commits January 20, 2025 11:43

ot.lp reorganise to avoid def in __init__

590e4d7

pr number + enabled pre-commit

109edb7

added barycenter.py imports

0957904

fixed wrong import in ot.gmm

818b3e7

ruff fix attempt

08c2285

removed ot bar contribs -> only o.lp reorganisation in this PR

f268515

add check_number_threads to ot/lp/__init__.py __all__

8f24cb9

update releases

3e3b444

made barycenter_solvers and network_simplex hidden + deprecated ot.lp…

566a0fc

….cvx

fix ref to lp.cvx in test

5c35d58

lp.cvx now imports barycenter and gives a warnings.warning

8ffb061

cvx import barycenter

26748eb

Merge branch 'PythonOT:master' into dev

d69bf97

added fixed-point barycenter function to ot.lp._barycenter_solvers_

081e4eb

github-actions bot added the ot.lp label Jan 20, 2025

ot bar demo

5952019

github-actions bot added the Examples label Jan 20, 2025

rflamary and others added 3 commits January 21, 2025 08:44

Merge branch 'master' into dev

6a3eab5

ot bar doc

3e8421e

doc fixes + ot bar coverage

ccf608a

github-actions bot added the Tests label Jan 21, 2025

python 3.13 in test workflow + added ggmot barycenter (WIP)

37b9c80

github-actions bot added the CI label Jan 21, 2025

eloitanguy added 4 commits January 21, 2025 16:06

fixed github action file

a20d3f0

ot bar doc + test coverage

0b6217b

examples: ot bar with projections onto circles + gmm ot bar

21bf86b

releases + readme + docs update

0820e51

eloitanguy changed the title ~~[WIP] OT barycenters for generic transport costs~~ [MRG] OT barycenters for generic transport costs Jan 21, 2025

Merge branch 'master' into dev

d1510ee

github-actions bot added the ot.gmm label Mar 3, 2025

eloitanguy and others added 2 commits March 12, 2025 15:11

Merge branch 'master' into dev

391ad39

ref fix

6bd4af8

rflamary reviewed Mar 14, 2025

View reviewed changes

eloitanguy and others added 2 commits March 17, 2025 19:54

implementation comments

51722bf

Merge branch 'master' into dev

371e3e7

eloitanguy changed the title ~~[MRG] OT barycenters for generic transport costs~~ [WIP] OT barycenters for generic transport costs Mar 18, 2025

rflamary and others added 8 commits March 25, 2025 15:59

Merge branch 'master' into dev

034daa0

Merge branch 'master' into dev

d81c9b6

Merge branch 'master' into dev

95e84ad

Merge branch 'master' into dev

869f58e

(WIP) added true barycenter fixed-point algorithm with updated tests …

f2269ac

…and examples

test and fixes

939b93d

no jax or tf support for free_support_generic_costs due to array assi…

6f47b29

…gnment

updated gmm bar colours

9a344e1

eloitanguy changed the title ~~[WIP] OT barycenters for generic transport costs~~ [MRG] OT barycenters for generic transport costs Jun 7, 2025

rflamary and others added 2 commits July 17, 2025 08:33

Merge branch 'master' into dev

d5ba13b

Merge branch 'master' into dev

56c1499

cedricvincentcuaz reviewed Aug 7, 2025

View reviewed changes


		torch.manual_seed(42)

		n = 136 # number of points of the of the barycentre

[MRG] OT barycenters for generic transport costs #715

Are you sure you want to change the base?

[MRG] OT barycenters for generic transport costs #715

Uh oh!

Conversation

eloitanguy commented Jan 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Types of changes

Motivation and context / Related issue

How has this been tested (if it applies)

PR checklist

Uh oh!

codecov bot commented Jan 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

rflamary left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

eloitanguy commented Mar 18, 2025

Uh oh!

eloitanguy commented Jun 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cedricvincentcuaz left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

eloitanguy commented Jan 20, 2025 •

edited

Loading

codecov bot commented Jan 20, 2025 •

edited

Loading

eloitanguy commented Jun 6, 2025 •

edited

Loading