1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
|
# -*- coding: utf-8 -*-
"""
LP solvers for optimal transport using cvxopt
"""
# Author: Remi Flamary <remi.flamary@unice.fr>
#
# License: MIT License
import numpy as np
import scipy as sp
import scipy.sparse as sps
try:
import cvxopt
from cvxopt import solvers, matrix, spmatrix
except ImportError:
cvxopt = False
def scipy_sparse_to_spmatrix(A):
"""Efficient conversion from scipy sparse matrix to cvxopt sparse matrix"""
coo = A.tocoo()
SP = spmatrix(coo.data.tolist(), coo.row.tolist(), coo.col.tolist(), size=A.shape)
return SP
def barycenter(A, M, weights=None, verbose=False, log=False, solver='interior-point'):
r"""Compute the Wasserstein barycenter of distributions A
The function solves the following optimization problem [16]:
.. math::
\mathbf{a} = arg\min_\mathbf{a} \sum_i W_{1}(\mathbf{a},\mathbf{a}_i)
where :
- :math:`W_1(\cdot,\cdot)` is the Wasserstein distance (see ot.emd.sinkhorn)
- :math:`\mathbf{a}_i` are training distributions in the columns of matrix :math:`\mathbf{A}`
The linear program is solved using the interior point solver from scipy.optimize.
If cvxopt solver if installed it can use cvxopt
Note that this problem do not scale well (both in memory and computational time).
Parameters
----------
A : np.ndarray (d,n)
n training distributions a_i of size d
M : np.ndarray (d,d)
loss matrix for OT
reg : float
Regularization term >0
weights : np.ndarray (n,)
Weights of each histogram a_i on the simplex (barycentric coodinates)
verbose : bool, optional
Print information along iterations
log : bool, optional
record log if True
solver : string, optional
the solver used, default 'interior-point' use the lp solver from
scipy.optimize. None, or 'glpk' or 'mosek' use the solver from cvxopt.
Returns
-------
a : (d,) ndarray
Wasserstein barycenter
log : dict
log dictionary return only if log==True in parameters
References
----------
.. [16] Agueh, M., & Carlier, G. (2011). Barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 43(2), 904-924.
"""
if weights is None:
weights = np.ones(A.shape[1]) / A.shape[1]
else:
assert(len(weights) == A.shape[1])
n_distributions = A.shape[1]
n = A.shape[0]
n2 = n * n
c = np.zeros((0))
b_eq1 = np.zeros((0))
for i in range(n_distributions):
c = np.concatenate((c, M.ravel() * weights[i]))
b_eq1 = np.concatenate((b_eq1, A[:, i]))
c = np.concatenate((c, np.zeros(n)))
lst_idiag1 = [sps.kron(sps.eye(n), np.ones((1, n))) for i in range(n_distributions)]
# row constraints
A_eq1 = sps.hstack((sps.block_diag(lst_idiag1), sps.coo_matrix((n_distributions * n, n))))
# columns constraints
lst_idiag2 = []
lst_eye = []
for i in range(n_distributions):
if i == 0:
lst_idiag2.append(sps.kron(np.ones((1, n)), sps.eye(n)))
lst_eye.append(-sps.eye(n))
else:
lst_idiag2.append(sps.kron(np.ones((1, n)), sps.eye(n - 1, n)))
lst_eye.append(-sps.eye(n - 1, n))
A_eq2 = sps.hstack((sps.block_diag(lst_idiag2), sps.vstack(lst_eye)))
b_eq2 = np.zeros((A_eq2.shape[0]))
# full problem
A_eq = sps.vstack((A_eq1, A_eq2))
b_eq = np.concatenate((b_eq1, b_eq2))
if not cvxopt or solver in ['interior-point']:
# cvxopt not installed or interior point
if solver is None:
solver = 'interior-point'
options = {'sparse': True, 'disp': verbose}
sol = sp.optimize.linprog(c, A_eq=A_eq, b_eq=b_eq, method=solver,
options=options)
x = sol.x
b = x[-n:]
else:
h = np.zeros((n_distributions * n2 + n))
G = -sps.eye(n_distributions * n2 + n)
sol = solvers.lp(matrix(c), scipy_sparse_to_spmatrix(G), matrix(h),
A=scipy_sparse_to_spmatrix(A_eq), b=matrix(b_eq),
solver=solver)
x = np.array(sol['x'])
b = x[-n:].ravel()
if log:
return b, sol
else:
return b
|
