Skip to content

Commit 23db72c

Browse files
Correct documentation for support barycenter (#201)
* example for log treatment in bregman.py * Improve doc * Revert "example for log treatment in bregman.py" This reverts commit 9f51c14 * Add comments by Flamary * Delete repetitive description * Added raw string to avoid pbs with backslashes
1 parent acfff52 commit 23db72c

File tree

1 file changed

+18
-8
lines changed

1 file changed

+18
-8
lines changed

ot/lp/__init__.py

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,7 @@ def emd(a, b, M, numItermax=100000, log=False, center_dual=True):
272272

273273
if np.any(~asel) or np.any(~bsel):
274274
u, v = estimate_dual_null_weights(u, v, a, b, M)
275-
275+
276276
result_code_string = check_result(result_code)
277277
if log:
278278
log = {}
@@ -389,7 +389,7 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(),
389389
if log or return_matrix:
390390
def f(b):
391391
bsel = b != 0
392-
392+
393393
G, cost, u, v, result_code = emd_c(a, b, M, numItermax)
394394

395395
if center_dual:
@@ -435,26 +435,36 @@ def f(b):
435435

436436
def free_support_barycenter(measures_locations, measures_weights, X_init, b=None, weights=None, numItermax=100,
437437
stopThr=1e-7, verbose=False, log=None):
438-
"""
439-
Solves the free support (locations of the barycenters are optimized, not the weights) Wasserstein barycenter problem (i.e. the weighted Frechet mean for the 2-Wasserstein distance)
438+
r"""
439+
Solves the free support (locations of the barycenters are optimized, not the weights) Wasserstein barycenter problem (i.e. the weighted Frechet mean for the 2-Wasserstein distance), formally:
440+
441+
.. math::
442+
\min_X \sum_{i=1}^N w_i W_2^2(b, X, a_i, X_i)
443+
444+
where :
445+
446+
- :math:`w \in \mathbb{(0, 1)}^{N}`'s are the barycenter weights and sum to one
447+
- the :math:`a_i \in \mathbb{R}^{k_i}` are the empirical measures weights and sum to one for each :math:`i`
448+
- the :math:`X_i \in \mathbb{R}^{k_i, d}` are the empirical measures atoms locations
449+
- :math:`b \in \mathbb{R}^{k}` is the desired weights vector of the barycenter
440450
441-
The function solves the Wasserstein barycenter problem when the barycenter measure is constrained to be supported on k atoms.
442451
This problem is considered in [1] (Algorithm 2). There are two differences with the following codes:
452+
443453
- we do not optimize over the weights
444454
- we do not do line search for the locations updates, we use i.e. theta = 1 in [1] (Algorithm 2). This can be seen as a discrete implementation of the fixed-point algorithm of [2] proposed in the continuous setting.
445455
446456
Parameters
447457
----------
448-
measures_locations : list of (k_i,d) numpy.ndarray
458+
measures_locations : list of N (k_i,d) numpy.ndarray
449459
The discrete support of a measure supported on k_i locations of a d-dimensional space (k_i can be different for each element of the list)
450-
measures_weights : list of (k_i,) numpy.ndarray
460+
measures_weights : list of N (k_i,) numpy.ndarray
451461
Numpy arrays where each numpy array has k_i non-negatives values summing to one representing the weights of each discrete input measure
452462
453463
X_init : (k,d) np.ndarray
454464
Initialization of the support locations (on k atoms) of the barycenter
455465
b : (k,) np.ndarray
456466
Initialization of the weights of the barycenter (non-negatives, sum to 1)
457-
weights : (k,) np.ndarray
467+
weights : (N,) np.ndarray
458468
Initialization of the coefficients of the barycenter (non-negatives, sum to 1)
459469
460470
numItermax : int, optional

0 commit comments

Comments
 (0)