Skip to content

[MRG] OT barycenters for generic transport costs #715

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 37 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
37 commits
Select commit Hold shift + click to select a range
590e4d7
ot.lp reorganise to avoid def in __init__
eloitanguy Jan 20, 2025
109edb7
pr number + enabled pre-commit
eloitanguy Jan 20, 2025
0957904
added barycenter.py imports
eloitanguy Jan 20, 2025
818b3e7
fixed wrong import in ot.gmm
eloitanguy Jan 20, 2025
08c2285
ruff fix attempt
eloitanguy Jan 20, 2025
f268515
removed ot bar contribs -> only o.lp reorganisation in this PR
eloitanguy Jan 20, 2025
8f24cb9
add check_number_threads to ot/lp/__init__.py __all__
eloitanguy Jan 20, 2025
3e3b444
update releases
eloitanguy Jan 20, 2025
566a0fc
made barycenter_solvers and network_simplex hidden + deprecated ot.lp…
eloitanguy Jan 20, 2025
5c35d58
fix ref to lp.cvx in test
eloitanguy Jan 20, 2025
8ffb061
lp.cvx now imports barycenter and gives a warnings.warning
eloitanguy Jan 20, 2025
26748eb
cvx import barycenter
eloitanguy Jan 20, 2025
d69bf97
Merge branch 'PythonOT:master' into dev
eloitanguy Jan 20, 2025
081e4eb
added fixed-point barycenter function to ot.lp._barycenter_solvers_
eloitanguy Jan 20, 2025
5952019
ot bar demo
eloitanguy Jan 20, 2025
6a3eab5
Merge branch 'master' into dev
rflamary Jan 21, 2025
3e8421e
ot bar doc
eloitanguy Jan 21, 2025
ccf608a
doc fixes + ot bar coverage
eloitanguy Jan 21, 2025
37b9c80
python 3.13 in test workflow + added ggmot barycenter (WIP)
eloitanguy Jan 21, 2025
a20d3f0
fixed github action file
eloitanguy Jan 21, 2025
0b6217b
ot bar doc + test coverage
eloitanguy Jan 21, 2025
21bf86b
examples: ot bar with projections onto circles + gmm ot bar
eloitanguy Jan 21, 2025
0820e51
releases + readme + docs update
eloitanguy Jan 21, 2025
d1510ee
Merge branch 'master' into dev
eloitanguy Mar 3, 2025
391ad39
Merge branch 'master' into dev
eloitanguy Mar 12, 2025
6bd4af8
ref fix
eloitanguy Mar 12, 2025
51722bf
implementation comments
eloitanguy Mar 17, 2025
371e3e7
Merge branch 'master' into dev
eloitanguy Mar 18, 2025
034daa0
Merge branch 'master' into dev
rflamary Mar 25, 2025
d81c9b6
Merge branch 'master' into dev
cedricvincentcuaz May 18, 2025
95e84ad
Merge branch 'master' into dev
rflamary May 23, 2025
869f58e
Merge branch 'master' into dev
eloitanguy Jun 5, 2025
f2269ac
(WIP) added true barycenter fixed-point algorithm with updated tests …
eloitanguy Jun 6, 2025
939b93d
test and fixes
eloitanguy Jun 6, 2025
6f47b29
no jax or tf support for free_support_generic_costs due to array assi…
eloitanguy Jun 6, 2025
9a344e1
updated gmm bar colours
eloitanguy Jun 6, 2025
d5ba13b
Merge branch 'master' into dev
rflamary Jul 17, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
python 3.13 in test workflow + added ggmot barycenter (WIP)
  • Loading branch information
eloitanguy committed Jan 21, 2025
commit 37b9c80cad43f3b71768a265a4c57ef57734e06c
2 changes: 1 addition & 1 deletion .github/workflows/build_tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ jobs:
strategy:
max-parallel: 4
matrix:
python-version: ["3.9", "3.10", "3.11", "3.12"]
python-version: ["3.9", "3.10", "3.11", "3.12, "3.13"]

steps:
- uses: actions/checkout@v4
Expand Down
114 changes: 113 additions & 1 deletion ot/gmm.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from .lp import emd2, emd
import numpy as np
from .utils import dist
from .gaussian import bures_wasserstein_mapping
from .gaussian import bures_wasserstein_mapping, bures_wasserstein_barycenter


def gaussian_logpdf(x, m, C):
Expand Down Expand Up @@ -440,3 +440,115 @@ def Tk0k1(k0, k1):
]
)
return nx.sum(mat, axis=(0, 1))


def solve_gmm_barycenter_fixed_point(
means,
covs,
means_list,
covs_list,
b_list,
weights,
max_its=300,
log=False,
barycentric_proj_method="euclidean",
):
r"""
Solves the GMM OT barycenter problem using the fixed point algorithm.

Parameters
----------
means : array-like
Initial (n, d) GMM means.
covs : array-like
Initial (n, d, d) GMM covariances.
means_list : list of array-like
List of K (m_k, d) GMM means.
covs_list : list of array-like
List of K (m_k, d, d) GMM covariances.
b_list : list of array-like
List of K (m_k) arrays of weights.
weights : array-like
Array (K,) of the barycentre coefficients.
max_its : int, optional
Maximum number of iterations (default is 300).
log : bool, optional
Whether to return the list of iterations (default is False).
barycentric_proj_method : str, optional
Method to project the barycentre weights: 'euclidean' (default) or 'bures'.

Returns
-------
means : array-like
(n, d) barycentre GMM means.
covs : array-like
(n, d, d) barycentre GMM covariances.
log_dict : dict, optional
Dictionary containing the list of iterations if log is True.
"""
nx = get_backend(means, covs[0], means_list[0], covs_list[0])
K = len(means_list)
n = means.shape[0]
d = means.shape[1]
means_its = [means.copy()]
covs_its = [covs.copy()]
a = nx.ones(n, type_as=means) / n

for _ in range(max_its):
pi_list = [
gmm_ot_plan(means, means_list[k], covs, covs_list[k], a, b_list[k])
for k in range(K)
]

means_selection, covs_selection = None, None
# in the euclidean case, the selection of Gaussians from each K sources
# comes from a barycentric projection is a convex combination of the
# selected means and covariances, which can be computed without a
# for loop on i
if barycentric_proj_method == "euclidean":
means_selection = nx.zeros((n, K, d), type_as=means)
covs_selection = nx.zeros((n, K, d, d), type_as=means)

for k in range(K):
means_selection[:, k, :] = n * pi_list[k] @ means_list[k]
covs_selection[:, k, :, :] = (
nx.einsum("ij,jab->iab", pi_list[k], covs_list[k]) * n
)

# each component i of the barycentre will be a Bures barycentre of the
# selected components of the K GMMs. In the 'bures' barycentric
# projection option, the selected components are also Bures barycentres.
for i in range(n):
# means_slice_i (K, d) is the selected means, each comes from a
# Gaussian barycentre along the disintegration of pi_k at i
# covs_slice_i (K, d, d) are the selected covariances
means_selection_i = []
covs_selection_i = []

# use previous computation (convex combination)
if barycentric_proj_method == "euclidean":
means_selection_i = means_selection[i]
covs_selection_i = covs_selection[i]

# compute Bures barycentre of the selected components
elif barycentric_proj_method == "bures":
w = (1 / a[i]) * pi_list[k][i, :]
for k in range(K):
m, C = bures_wasserstein_barycenter(means_list[k], covs_list[k], w)
means_selection_i.append(m)
covs_selection_i.append(C)

else:
raise ValueError("Unknown barycentric_proj_method")

means[i], covs[i] = bures_wasserstein_barycenter(
means_selection_i, covs_selection_i, weights
)

if log:
means_its.append(means.copy())
covs_its.append(covs.copy())

if log:
return means, covs, {"means_its": means_its, "covs_its": covs_its}
return means, covs
Loading