# -*- coding: utf-8 -*-
"""
Regularization path OT solvers
"""
# Author: Haoran Wu <haoran.wu@univ-ubs.fr>
# License: MIT License
import numpy as np
import scipy.sparse as sp
[docs]
def recast_ot_as_lasso(a, b, C):
r"""This function recasts the l2-penalized UOT problem as a Lasso problem.
Recall the l2-penalized UOT problem defined in
:ref:`[41] <references-regpath>`
.. math::
\text{UOT}_{\lambda} = \min_T <C, T> + \lambda \|T 1_m -
\mathbf{a}\|_2^2 +
\lambda \|T^T 1_n - \mathbf{b}\|_2^2
s.t.
T \geq 0
where :
- :math:`C` is the cost matrix
- :math:`\lambda` is the l2-regularization parameter
- :math:`\mathbf{a}` and :math:`\mathbf{b}` are the source and target \
distributions
- :math:`T` is the transport plan to optimize
The problem above can be reformulated as a non-negative penalized
linear regression problem, particularly Lasso
.. math::
\text{UOT2}_{\lambda} = \min_{\mathbf{t}} \gamma \mathbf{c}^T
\mathbf{t} + 0.5 * \|H \mathbf{t} - \mathbf{y}\|_2^2
s.t.
\mathbf{t} \geq 0
where :
- :math:`\mathbf{c}` is the flattened version of the cost matrix :math:`C`
- :math:`\mathbf{y}` is the concatenation of vectors :math:`\mathbf{a}` \
and :math:`\mathbf{b}`
- :math:`H` is a metric matrix, see :ref:`[41] <references-regpath>` for \
the design of :math:`H`. The matrix product :math:`H\mathbf{t}` \
computes both the source marginal and the target marginals.
- :math:`\mathbf{t}` is the flattened version of the transport plan \
:math:`T`
Parameters
----------
a : np.ndarray (dim_a,)
Histogram of dimension dim_a
b : np.ndarray (dim_b,)
Histogram of dimension dim_b
C : np.ndarray, shape (dim_a, dim_b)
Cost matrix
Returns
-------
H : np.ndarray (dim_a+dim_b, dim_a*dim_b)
Design matrix that contains only 0 and 1
y : np.ndarray (ns + nt, )
Concatenation of histograms :math:`\mathbf{a}` and :math:`\mathbf{b}`
c : np.ndarray (ns * nt, )
Flattened array of the cost matrix
Examples
--------
>>> import ot
>>> a = np.array([0.2, 0.3, 0.5])
>>> b = np.array([0.1, 0.9])
>>> C = np.array([[16., 25.], [28., 16.], [40., 36.]])
>>> H, y, c = ot.regpath.recast_ot_as_lasso(a, b, C)
>>> H.toarray()
array([[1., 1., 0., 0., 0., 0.],
[0., 0., 1., 1., 0., 0.],
[0., 0., 0., 0., 1., 1.],
[1., 0., 1., 0., 1., 0.],
[0., 1., 0., 1., 0., 1.]])
>>> y
array([0.2, 0.3, 0.5, 0.1, 0.9])
>>> c
array([16., 25., 28., 16., 40., 36.])
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
dim_a = np.shape(a)[0]
dim_b = np.shape(b)[0]
y = np.concatenate((a, b))
c = C.flatten()
jHa = np.arange(dim_a * dim_b)
iHa = np.repeat(np.arange(dim_a), dim_b)
jHb = np.arange(dim_a * dim_b)
iHb = np.tile(np.arange(dim_b), dim_a) + dim_a
j = np.concatenate((jHa, jHb))
i = np.concatenate((iHa, iHb))
H = sp.csc_matrix(
(np.ones(dim_a * dim_b * 2), (i, j)), shape=(dim_a + dim_b, dim_a * dim_b)
)
return H, y, c
[docs]
def recast_semi_relaxed_as_lasso(a, b, C):
r"""This function recasts the semi-relaxed l2-UOT problem as Lasso problem.
.. math::
\text{semi-relaxed UOT} = \min_T <C, T>
+ \lambda \|T 1_m - \mathbf{a}\|_2^2
s.t.
T^T 1_n = \mathbf{b}
\mathbf{t} \geq 0
where :
- :math:`C` is the metric cost matrix
- :math:`\lambda` is the l2-regularization parameter
- :math:`\mathbf{a}` and :math:`\mathbf{b}` are the source and target \
distributions
- :math:`T` is the transport plan to optimize
The problem above can be reformulated as follows
.. math::
\text{semi-relaxed UOT2} = \min_t \gamma \mathbf{c}^T t
+ 0.5 * \|H_r \mathbf{t} - \mathbf{a}\|_2^2
s.t.
H_c \mathbf{t} = \mathbf{b}
\mathbf{t} \geq 0
where :
- :math:`\mathbf{c}` is flattened version of the cost matrix :math:`C`
- :math:`\gamma = 1/\lambda` is the l2-regularization parameter
- :math:`H_r` is a metric matrix which computes the sum along the \
rows of the transport plan :math:`T`
- :math:`H_c` is a metric matrix which computes the sum along the \
columns of the transport plan :math:`T`
- :math:`\mathbf{t}` is the flattened version of :math:`T`
Parameters
----------
a : np.ndarray (dim_a,)
Histogram of dimension dim_a
b : np.ndarray (dim_b,)
Histogram of dimension dim_b
C : np.ndarray, shape (dim_a, dim_b)
Cost matrix
Returns
-------
Hr : np.ndarray (dim_a, dim_a * dim_b)
Auxiliary matrix constituted by 0 and 1, which computes
the sum along the rows of transport plan :math:`T`
Hc : np.ndarray (dim_b, dim_a * dim_b)
Auxiliary matrix constituted by 0 and 1, which computes
the sum along the columns of transport plan :math:`T`
c : np.ndarray (ns * nt, )
Flattened array of the cost matrix
Examples
--------
>>> import ot
>>> a = np.array([0.2, 0.3, 0.5])
>>> b = np.array([0.1, 0.9])
>>> C = np.array([[16., 25.], [28., 16.], [40., 36.]])
>>> Hr,Hc,c = ot.regpath.recast_semi_relaxed_as_lasso(a, b, C)
>>> Hr.toarray()
array([[1., 1., 0., 0., 0., 0.],
[0., 0., 1., 1., 0., 0.],
[0., 0., 0., 0., 1., 1.]])
>>> Hc.toarray()
array([[1., 0., 1., 0., 1., 0.],
[0., 1., 0., 1., 0., 1.]])
>>> c
array([16., 25., 28., 16., 40., 36.])
"""
dim_a = np.shape(a)[0]
dim_b = np.shape(b)[0]
c = C.flatten()
jHr = np.arange(dim_a * dim_b)
iHr = np.repeat(np.arange(dim_a), dim_b)
jHc = np.arange(dim_a * dim_b)
iHc = np.tile(np.arange(dim_b), dim_a)
Hr = sp.csc_matrix(
(np.ones(dim_a * dim_b), (iHr, jHr)), shape=(dim_a, dim_a * dim_b)
)
Hc = sp.csc_matrix(
(np.ones(dim_a * dim_b), (iHc, jHc)), shape=(dim_b, dim_a * dim_b)
)
return Hr, Hc, c
[docs]
def ot_next_gamma(phi, delta, HtH, Hty, c, active_index, current_gamma):
r""" This function computes the next value of gamma if a variable
is added in the next iteration of the regularization path.
We look for the largest value of gamma such that
the gradient of an inactive variable vanishes
.. math::
\max_{i \in \bar{A}} \frac{\mathbf{h}_i^T(H_A \phi - \mathbf{y})}
{\mathbf{h}_i^T H_A \delta - \mathbf{c}_i}
where :
- A is the current active set
- :math:`\mathbf{h}_i` is the :math:`i` th column of the design \
matrix :math:`{H}`
- :math:`{H}_A` is the sub-matrix constructed by the columns of \
:math:`{H}` whose indices belong to the active set A
- :math:`\mathbf{c}_i` is the :math:`i` th element of the cost vector \
:math:`\mathbf{c}`
- :math:`\mathbf{y}` is the concatenation of the source and target \
distributions
- :math:`\phi` is the intercept of the solutions at the current iteration
- :math:`\delta` is the slope of the solutions at the current iteration
Parameters
----------
phi : np.ndarray (size(A), )
Intercept of the solutions at the current iteration
delta : np.ndarray (size(A), )
Slope of the solutions at the current iteration
HtH : np.ndarray (dim_a * dim_b, dim_a * dim_b)
Matrix product of :math:`{H}^T {H}`
Hty : np.ndarray (dim_a + dim_b, )
Matrix product of :math:`{H}^T \mathbf{y}`
c: np.ndarray (dim_a * dim_b, )
Flattened array of the cost matrix :math:`{C}`
active_index : list
Indices of active variables
current_gamma : float
Value of the regularization parameter at the beginning of the current \
iteration
Returns
-------
next_gamma : float
Value of gamma if a variable is added to active set in next iteration
next_active_index : int
Index of variable to be activated
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
M = (HtH[:, active_index].dot(phi) - Hty) / (
HtH[:, active_index].dot(delta) - c + 1e-16
)
M[active_index] = 0
M[M > (current_gamma - 1e-10 * current_gamma)] = 0
return np.max(M), np.argmax(M)
[docs]
def semi_relaxed_next_gamma(
phi, delta, phi_u, delta_u, HrHr, Hc, Hra, c, active_index, current_gamma
):
r""" This function computes the next value of gamma when a variable is
active in the regularization path of semi-relaxed UOT.
By taking the Lagrangian form of the problem, we obtain a similar update
as the two-sided relaxed UOT
.. math::
\max_{i \in \bar{A}} \frac{\mathbf{h}_{ri}^T(H_{rA} \phi - \mathbf{a})
+ \mathbf{h}_{c i}^T\phi_u}{\mathbf{h}_{r i}^T H_{r A} \delta + \
\mathbf{h}_{c i} \delta_u - \mathbf{c}_i}
where :
- A is the current active set
- :math:`\mathbf{h}_{r i}` is the ith column of the matrix :math:`H_r`
- :math:`\mathbf{h}_{c i}` is the ith column of the matrix :math:`H_c`
- :math:`H_{r A}` is the sub-matrix constructed by the columns of \
:math:`H_r` whose indices belong to the active set A
- :math:`\mathbf{c}_i` is the :math:`i` th element of cost vector \
:math:`\mathbf{c}`
- :math:`\phi` is the intercept of the solutions in current iteration
- :math:`\delta` is the slope of the solutions in current iteration
- :math:`\phi_u` is the intercept of Lagrange parameter at the \
current iteration
- :math:`\delta_u` is the slope of Lagrange parameter at the \
current iteration
Parameters
----------
phi : np.ndarray (size(A), )
Intercept of the solutions at the current iteration
delta : np.ndarray (size(A), )
Slope of the solutions at the current iteration
phi_u : np.ndarray (dim_b, )
Intercept of the Lagrange parameter at the current iteration
delta_u : np.ndarray (dim_b, )
Slope of the Lagrange parameter at the current iteration
HrHr : np.ndarray (dim_a * dim_b, dim_a * dim_b)
Matrix product of :math:`H_r^T H_r`
Hc : np.ndarray (dim_b, dim_a * dim_b)
Matrix that computes the sum along the columns of the transport plan \
:math:`T`
Hra : np.ndarray (dim_a * dim_b, )
Matrix product of :math:`H_r^T \mathbf{a}`
c: np.ndarray (dim_a * dim_b, )
Flattened array of cost matrix :math:`C`
active_index : list
Indices of active variables
current_gamma : float
Value of regularization coefficient at the start of current iteration
Returns
-------
next_gamma : float
Value of gamma if a variable is added to active set in next iteration
next_active_index : int
Index of variable to be activated
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
M = (HrHr[:, active_index].dot(phi) - Hra + Hc.T.dot(phi_u)) / (
HrHr[:, active_index].dot(delta) - c + Hc.T.dot(delta_u) + 1e-16
)
M[active_index] = 0
M[M > (current_gamma - 1e-10 * current_gamma)] = 0
return np.max(M), np.argmax(M)
[docs]
def compute_next_removal(phi, delta, current_gamma):
r""" This function computes the next gamma value if a variable
is removed at the next iteration of the regularization path.
We look for the largest value of the regularization parameter such that
an element of the current solution vanishes
.. math::
\max_{j \in A} \frac{\phi_j}{\delta_j}
where :
- A is the current active set
- :math:`\phi_j` is the :math:`j` th element of the intercept of the \
current solution
- :math:`\delta_j` is the :math:`j` th element of the slope of the \
current solution
Parameters
----------
phi : ndarray, shape (size(A), )
Intercept of the solution at the current iteration
delta : ndarray, shape (size(A), )
Slope of the solution at the current iteration
current_gamma : float
Value of the regularization parameter at the beginning of the \
current iteration
Returns
-------
next_removal_gamma : float
Gamma value if a variable is removed at the next iteration
next_removal_index : int
Index of the variable to be removed at the next iteration
.. _references-regpath:
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
r_candidate = phi / (delta - 1e-16)
r_candidate[r_candidate >= (1 - 1e-8) * current_gamma] = 0
return np.max(r_candidate), np.argmax(r_candidate)
[docs]
def complement_schur(M_current, b, d, id_pop):
r""" This function computes the inverse of the design matrix in the \
regularization path using the Schur complement. Two cases may arise:
Case 1: one variable is added to the active set
.. math::
M_{k+1}^{-1} =
\begin{bmatrix}
M_{k}^{-1} + s^{-1} M_{k}^{-1} \mathbf{b} \mathbf{b}^T M_{k}^{-1} \
& - M_{k}^{-1} \mathbf{b} s^{-1} \\
- s^{-1} \mathbf{b}^T M_{k}^{-1} & s^{-1}
\end{bmatrix}
where :
- :math:`M_k^{-1}` is the inverse of the design matrix :math:`H_A^tH_A` \
of the previous iteration
- :math:`\mathbf{b}` is the last column of :math:`M_{k}`
- :math:`s` is the Schur complement, given by \
:math:`s = \mathbf{d} - \mathbf{b}^T M_{k}^{-1} \mathbf{b}`
Case 2: one variable is removed from the active set.
.. math::
M_{k+1}^{-1} = M^{-1}_{k \backslash q} -
\frac{r_{-q,q} r^{T}_{-q,q}}{r_{q,q}}
where :
- :math:`q` is the index of column and row to delete
- :math:`M^{-1}_{k \backslash q}` is the previous inverse matrix deprived \
of the :math:`q` th column and :math:`q` th row
- :math:`r_{-q,q}` is the :math:`q` th column of :math:`M^{-1}_{k}` \
without the :math:`q` th element
- :math:`r_{q, q}` is the element of :math:`q` th column and :math:`q` th \
row in :math:`M^{-1}_{k}`
Parameters
----------
M_current : ndarray, shape (size(A)-1, size(A)-1)
Inverse matrix of :math:`H_A^tH_A` at the previous iteration, with \
size(A) the size of the active set
b : ndarray, shape (size(A)-1, )
None for case 2 (removal), last column of :math:`M_{k}` for case 1 \
(addition)
d : float
should be equal to 2 when UOT and 1 for the semi-relaxed OT
id_pop : int
Index of the variable to be removed, equal to -1
if no variable is deleted at the current iteration
Returns
-------
M : ndarray, shape (size(A), size(A))
Inverse matrix of :math:`H_A^tH_A` of the current iteration
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
if b is None:
b = M_current[id_pop, :]
b = np.delete(b, id_pop)
M_del = np.delete(M_current, id_pop, 0)
a = M_del[:, id_pop]
M_del = np.delete(M_del, id_pop, 1)
M = M_del - np.outer(a, b) / M_current[id_pop, id_pop]
else:
n = b.shape[0] + 1
if np.shape(b)[0] == 0:
M = np.array([[0.5]])
else:
X = M_current.dot(b)
s = d - b.T.dot(X)
M = np.zeros((n, n))
M[:-1, :-1] = M_current + X.dot(X.T) / s
X_ravel = X.ravel()
M[-1, :-1] = -X_ravel / s
M[:-1, -1] = -X_ravel / s
M[-1, -1] = 1 / s
return M
[docs]
def construct_augmented_H(active_index, m, Hc, HrHr):
r""" This function constructs an augmented matrix for the first iteration
of the semi-relaxed regularization path
.. math::
\text{Augmented}_H =
\begin{bmatrix}
0 & H_{c A} \\
H_{c A}^T & H_{r A}^T H_{r A}
\end{bmatrix}
where :
- :math:`H_{r A}` is the sub-matrix constructed by the columns of \
:math:`H_r` whose indices belong to the active set A
- :math:`H_{c A}` is the sub-matrix constructed by the columns of \
:math:`H_c` whose indices belong to the active set A
Parameters
----------
active_index : list
Indices of the active variables
m : int
Length of the target distribution
Hc : np.ndarray (dim_b, dim_a * dim_b)
Matrix that computes the sum along the columns of the transport plan \
:math:`T`
HrHr : np.ndarray (dim_a * dim_b, dim_a * dim_b)
Matrix product of :math:`H_r^T H_r`
Returns
-------
H_augmented : np.ndarray (dim_b + size(A), dim_b + size(A))
Augmented matrix for the first iteration of the semi-relaxed
regularization path
"""
Hc_sub = Hc[:, active_index].toarray()
HrHr_sub = HrHr[:, active_index]
HrHr_sub = HrHr_sub[active_index, :].toarray()
H_augmented = np.block([[np.zeros((m, m)), Hc_sub], [Hc_sub.T, HrHr_sub]])
return H_augmented
[docs]
def fully_relaxed_path(a: np.array, b: np.array, C: np.array, reg=1e-4, itmax=50000):
r"""This function gives the regularization path of l2-penalized UOT problem
The problem to optimize is the Lasso reformulation of the l2-penalized UOT:
.. math::
\min_t \gamma \mathbf{c}^T \mathbf{t}
+ 0.5 * \|{H} \mathbf{t} - \mathbf{y}\|_2^2
s.t.
\mathbf{t} \geq 0
where :
- :math:`\mathbf{c}` is the flattened version of the cost matrix \
:math:`{C}`
- :math:`\gamma = 1/\lambda` is the l2-regularization coefficient
- :math:`\mathbf{y}` is the concatenation of vectors :math:`\mathbf{a}` \
and :math:`\mathbf{b}`, defined as \
:math:`\mathbf{y}^T = [\mathbf{a}^T \mathbf{b}^T]`
- :math:`{H}` is a design matrix, see :ref:`[41] <references-regpath>` \
for the design of :math:`{H}`. The matrix product :math:`H\mathbf{t}` \
computes both the source marginal and the target marginals.
- :math:`\mathbf{t}` is the flattened version of the transport matrix
Parameters
----------
a : np.ndarray (dim_a,)
Histogram of dimension dim_a
b : np.ndarray (dim_b,)
Histogram of dimension dim_b
C : np.ndarray, shape (dim_a, dim_b)
Cost matrix
reg: float
l2-regularization coefficient
itmax: int
Maximum number of iteration
Returns
-------
t : np.ndarray (dim_a*dim_b, )
Flattened vector of the optimal transport matrix
t_list : list
List of solutions in the regularization path
gamma_list : list
List of regularization coefficients in the regularization path
Examples
--------
>>> import ot
>>> import numpy as np
>>> n = 3
>>> xs = np.array([1., 2., 3.]).reshape((n, 1))
>>> xt = np.array([5., 6., 7.]).reshape((n, 1))
>>> C = ot.dist(xs, xt)
>>> C /= C.max()
>>> a = np.array([0.2, 0.5, 0.3])
>>> b = np.array([0.2, 0.5, 0.3])
>>> t, _, _ = ot.regpath.fully_relaxed_path(a, b, C, 1e-4)
>>> t
array([1.99958333e-01, 0.00000000e+00, 0.00000000e+00, 3.88888889e-05,
4.99938889e-01, 0.00000000e+00, 0.00000000e+00, 3.88888889e-05,
2.99958333e-01])
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
n = np.shape(a)[0]
m = np.shape(b)[0]
H, y, c = recast_ot_as_lasso(a, b, C)
HtH = H.T.dot(H)
Hty = H.T.dot(y)
n_iter = 1
# initialization
M0 = Hty / c
gamma_list = [np.max(M0)]
active_index = [np.argmax(M0)]
t_list = [np.zeros((n * m,))]
H_inv = np.array([[]])
add_col = np.array([])
id_pop = -1
while n_iter < itmax and gamma_list[-1] > reg:
H_inv = complement_schur(H_inv, add_col, 2.0, id_pop)
current_gamma = gamma_list[-1]
# compute the intercept and slope of solutions in current iteration
# t = phi - gamma * delta
phi = H_inv.dot(Hty[active_index])
delta = H_inv.dot(c[active_index])
gamma, ik = ot_next_gamma(phi, delta, HtH, Hty, c, active_index, current_gamma)
# compute the next lambda when removing a point from the active set
alt_gamma, id_pop = compute_next_removal(phi, delta, current_gamma)
# if the positivity constraint is violated, we remove id_pop
# from active set, otherwise we add ik to active set
if alt_gamma > gamma:
gamma = alt_gamma
else:
id_pop = -1
# compute the solution of current segment
tA = phi - gamma * delta
sol = np.zeros((n * m,))
sol[active_index] = tA
if id_pop != -1:
active_index.pop(id_pop)
add_col = None
else:
active_index.append(ik)
add_col = HtH[active_index[:-1], ik].toarray()
gamma_list.append(gamma)
t_list.append(sol)
n_iter += 1
if itmax <= n_iter:
print("maximum iteration has been reached !")
# correct the last solution and gamma
if len(t_list) > 1:
t_final = t_list[-2] + (t_list[-1] - t_list[-2]) * (reg - gamma_list[-2]) / (
gamma_list[-1] - gamma_list[-2]
)
t_list[-1] = t_final
gamma_list[-1] = reg
else:
gamma_list[-1] = reg
print("Regularization path does not exist !")
return t_list[-1], t_list, gamma_list
[docs]
def semi_relaxed_path(a: np.array, b: np.array, C: np.array, reg=1e-4, itmax=50000):
r"""This function gives the regularization path of semi-relaxed
l2-UOT problem.
The problem to optimize is the Lasso reformulation of the l2-penalized UOT:
.. math::
\min_t \gamma \mathbf{c}^T t
+ 0.5 * \|H_r \mathbf{t} - \mathbf{a}\|_2^2
s.t.
H_c \mathbf{t} = \mathbf{b}
\mathbf{t} \geq 0
where :
- :math:`\mathbf{c}` is the flattened version of the cost matrix \
:math:`C`
- :math:`\gamma = 1/\lambda` is the l2-regularization parameter
- :math:`H_r` is a matrix that computes the sum along the rows of \
the transport plan :math:`T`
- :math:`H_c` is a matrix that computes the sum along the columns of \
the transport plan :math:`T`
- :math:`\mathbf{t}` is the flattened version of the transport plan \
:math:`T`
Parameters
----------
a : np.ndarray (dim_a,)
Histogram of dimension dim_a
b : np.ndarray (dim_b,)
Histogram of dimension dim_b
C : np.ndarray, shape (dim_a, dim_b)
Cost matrix
reg: float (optional)
l2-regularization coefficient
itmax: int (optional)
Maximum number of iteration
Returns
-------
t : np.ndarray (dim_a*dim_b, )
Flattened vector of the (unregularized) optimal transport matrix
t_list : list
List of all the optimal transport vectors of the regularization path
gamma_list : list
List of the regularization parameters in the path
Examples
--------
>>> import ot
>>> import numpy as np
>>> n = 3
>>> xs = np.array([1., 2., 3.]).reshape((n, 1))
>>> xt = np.array([5., 6., 7.]).reshape((n, 1))
>>> C = ot.dist(xs, xt)
>>> C /= C.max()
>>> a = np.array([0.2, 0.5, 0.3])
>>> b = np.array([0.2, 0.5, 0.3])
>>> t, _, _ = ot.regpath.semi_relaxed_path(a, b, C, 1e-4)
>>> t
array([1.99980556e-01, 0.00000000e+00, 0.00000000e+00, 1.94444444e-05,
4.99980556e-01, 0.00000000e+00, 0.00000000e+00, 1.94444444e-05,
3.00000000e-01])
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
n = np.shape(a)[0]
m = np.shape(b)[0]
Hr, Hc, c = recast_semi_relaxed_as_lasso(a, b, C)
Hra = Hr.T.dot(a)
HrHr = Hr.T.dot(Hr)
n_iter = 1
active_index = []
# initialization
for j in range(np.shape(C)[1]):
i = np.argmin(C[:, j])
active_index.append(i * m + j)
gamma_list = []
t_list = []
current_gamma = np.inf
augmented_H0 = construct_augmented_H(active_index, m, Hc, HrHr)
add_col = np.array([])
id_pop = -1
while n_iter < itmax and current_gamma > reg:
if n_iter == 1:
H_inv = np.linalg.inv(augmented_H0)
else:
H_inv = complement_schur(H_inv, add_col, 1.0, id_pop + m)
# compute the intercept and slope of solutions in current iteration
augmented_phi = H_inv.dot(np.concatenate((b, Hra[active_index])))
augmented_delta = H_inv[:, m:].dot(c[active_index])
phi = augmented_phi[m:]
delta = augmented_delta[m:]
phi_u = augmented_phi[0:m]
delta_u = augmented_delta[0:m]
gamma, ik = semi_relaxed_next_gamma(
phi, delta, phi_u, delta_u, HrHr, Hc, Hra, c, active_index, current_gamma
)
# compute the next lambda when removing a point from the active set
alt_gamma, id_pop = compute_next_removal(phi, delta, current_gamma)
# if the positivity constraint is violated, we remove id_pop
# from active set, otherwise we add ik to active set
if alt_gamma > gamma:
gamma = alt_gamma
else:
id_pop = -1
# compute the solution of current segment
tA = phi - gamma * delta
sol = np.zeros((n * m,))
sol[active_index] = tA
if id_pop != -1:
active_index.pop(id_pop)
add_col = None
else:
active_index.append(ik)
add_col = np.concatenate(
(Hc.toarray()[:, ik], HrHr.toarray()[active_index[:-1], ik])
)
add_col = add_col[:, np.newaxis]
gamma_list.append(gamma)
t_list.append(sol)
current_gamma = gamma
n_iter += 1
if itmax <= n_iter:
print("maximum iteration has been reached !")
# correct the last solution and gamma
if len(t_list) > 1:
t_final = t_list[-2] + (t_list[-1] - t_list[-2]) * (reg - gamma_list[-2]) / (
gamma_list[-1] - gamma_list[-2]
)
t_list[-1] = t_final
gamma_list[-1] = reg
else:
gamma_list[-1] = reg
print("Regularization path does not exist !")
return t_list[-1], t_list, gamma_list
[docs]
def regularization_path(
a: np.array, b: np.array, C: np.array, reg=1e-4, semi_relaxed=False, itmax=50000
):
r"""This function provides all the solutions of the regularization path \
of the l2-UOT problem :ref:`[41] <references-regpath>`.
The problem to optimize is the Lasso reformulation of the l2-penalized UOT:
.. math::
\min_t \gamma \mathbf{c}^T \mathbf{t}
+ 0.5 * \|{H} \mathbf{t} - \mathbf{y}\|_2^2
s.t.
\mathbf{t} \geq 0
where :
- :math:`\mathbf{c}` is the flattened version of the cost matrix \
:math:`{C}`
- :math:`\gamma = 1/\lambda` is the l2-regularization coefficient
- :math:`\mathbf{y}` is the concatenation of vectors :math:`\mathbf{a}` \
and :math:`\mathbf{b}`, defined as \
:math:`\mathbf{y}^T = [\mathbf{a}^T \mathbf{b}^T]`
- :math:`{H}` is a design matrix, see :ref:`[41] <references-regpath>` \
for the design of :math:`{H}`. The matrix product :math:`H\mathbf{t}` \
computes both the source marginal and the target marginals.
- :math:`\mathbf{t}` is the flattened version of the transport matrix
For the semi-relaxed problem, it optimizes the Lasso reformulation of the
l2-penalized UOT:
.. math::
\min_t \gamma \mathbf{c}^T \mathbf{t}
+ 0.5 * \|H_r \mathbf{t} - \mathbf{a}\|_2^2
s.t.
H_c \mathbf{t} = \mathbf{b}
\mathbf{t} \geq 0
Parameters
----------
a : np.ndarray (dim_a,)
Histogram of dimension dim_a
b : np.ndarray (dim_b,)
Histogram of dimension dim_b
C : np.ndarray, shape (dim_a, dim_b)
Cost matrix
reg: float (optional)
l2-regularization coefficient
semi_relaxed : bool (optional)
Give the semi-relaxed path if True
itmax: int (optional)
Maximum number of iteration
Returns
-------
t : np.ndarray (dim_a*dim_b, )
Flattened vector of the (unregularized) optimal transport matrix
t_list : list
List of all the optimal transport vectors of the regularization path
gamma_list : list
List of the regularization parameters in the path
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
if semi_relaxed:
t, t_list, gamma_list = semi_relaxed_path(a, b, C, reg=reg, itmax=itmax)
else:
t, t_list, gamma_list = fully_relaxed_path(a, b, C, reg=reg, itmax=itmax)
return t, t_list, gamma_list
[docs]
def compute_transport_plan(gamma, gamma_list, Pi_list):
r"""Given the regularization path, this function computes the transport
plan for any value of gamma thanks to the piecewise linearity of the path.
.. math::
t(\gamma) = \phi(\gamma) - \gamma \delta(\gamma)
where:
- :math:`\gamma` is the regularization parameter
- :math:`\phi(\gamma)` is the corresponding intercept
- :math:`\delta(\gamma)` is the corresponding slope
- :math:`\mathbf{t}` is the flattened version of the transport matrix
Parameters
----------
gamma : float
Regularization coefficient
gamma_list : list
List of regularization parameters of the regularization path
Pi_list : list
List of all the solutions of the regularization path
Returns
-------
t : np.ndarray (dim_a*dim_b, )
Vectorization of the transport plan corresponding to the given value
of gamma
Examples
--------
>>> import ot
>>> import numpy as np
>>> n = 3
>>> xs = np.array([1., 2., 3.]).reshape((n, 1))
>>> xt = np.array([5., 6., 7.]).reshape((n, 1))
>>> C = ot.dist(xs, xt)
>>> C /= C.max()
>>> a = np.array([0.2, 0.5, 0.3])
>>> b = np.array([0.2, 0.5, 0.3])
>>> t, pi_list, g_list = ot.regpath.regularization_path(a, b, C, reg=1e-4)
>>> gamma = 1
>>> t2 = ot.regpath.compute_transport_plan(gamma, g_list, pi_list)
>>> t2
array([0. , 0. , 0. , 0.19722222, 0.05555556,
0. , 0. , 0.24722222, 0. ])
.. _references-regpath:
References
----------
.. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
Unbalanced optimal transport through non-negative penalized
linear regression. NeurIPS.
"""
if gamma >= gamma_list[0]:
Pi = Pi_list[0]
elif gamma <= gamma_list[-1]:
Pi = Pi_list[-1]
else:
idx = np.where(gamma <= np.array(gamma_list))[0][-1]
gamma_k0 = gamma_list[idx]
gamma_k1 = gamma_list[idx + 1]
pi_k0 = Pi_list[idx]
pi_k1 = Pi_list[idx + 1]
Pi = pi_k0 + (pi_k1 - pi_k0) * (gamma - gamma_k0) / (gamma_k1 - gamma_k0)
return Pi