diff options
author | haoran010 <62598274+haoran010@users.noreply.github.com> | 2021-10-25 10:47:22 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2021-10-25 10:47:22 +0200 |
commit | 7af8c2147d61349f4d99ca33318a8a125e4569aa (patch) | |
tree | 5c08a89f2c998a6c1d734be28e4127130c1c1102 /ot/regpath.py | |
parent | d50d8145a5c0cf69d438b018cd5f1b914905e784 (diff) |
[MRG] Regularization path for l2 UOT (#274)
* add reg path
* debug examples and verify pep8
* pep8 and move the reg path examples in unbalanced folder
Co-authored-by: haoran010 <haoran.wu@insa-rennes.fr>
Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
Diffstat (limited to 'ot/regpath.py')
-rw-r--r-- | ot/regpath.py | 827 |
1 files changed, 827 insertions, 0 deletions
diff --git a/ot/regpath.py b/ot/regpath.py new file mode 100644 index 0000000..269937a --- /dev/null +++ b/ot/regpath.py @@ -0,0 +1,827 @@ +# -*- 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 + + +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 [Chapel et al., 2021] + .. math:: + UOT = \min_T <C, T> + \lambda \|T 1_m - a\|_2^2 + + \lambda \|T^T 1_n - b\|_2^2 + s.t. + T \geq 0 + where : + - C is the (dim_a, dim_b) metric cost matrix + - :math:`\lambda` is the l2-regularization coefficient + - a and b are source and target distributions + - T is the transport plan to optimize + + The problem above can be reformulated to a non-negative penalized + linear regression problem, particularly Lasso + .. math:: + UOT2 = \min_t \gamma c^T t + 0.5 * \|H t - y\|_2^2 + s.t. + t \geq 0 + where : + - c is a (dim_a * dim_b, ) metric cost vector (flattened version of C) + - :math:`\gamma = 1/\lambda` is the l2-regularization coefficient + - y is the concatenation of vectors a and b, defined as y^T = [a^T b^T] + - H is a (dim_a + dim_b, dim_a * dim_b) metric matrix, + see [Chapel et al., 2021] for the design of H. The matrix product H t + computes both the source marginal and the target marginal. + - t is a (dim_a * dim_b, ) metric vector (flattened version of 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) + Auxiliary matrix constituted by 0 and 1 + y : np.ndarray (ns + nt, ) + Concatenation of histogram a and histogram b + c : np.ndarray (ns * nt, ) + Flattened array of 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 + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + + 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 + + +def recast_semi_relaxed_as_lasso(a, b, C): + r"""This function recasts the semi-relaxed l2-UOT problem as Lasso problem + + .. math:: + semi-relaxed UOT = \min_T <C, T> + \lambda \|T 1_m - a\|_2^2 + s.t. + T^T 1_n = b + t \geq 0 + where : + - C is the (dim_a, dim_b) metric cost matrix + - :math:`\lambda` is the l2-regularization coefficient + - a and b are source and target distributions + - T is the transport plan to optimize + + The problem above can be reformulated as follows + .. math:: + semi-relaxed UOT2 = \min_t \gamma c^T t + 0.5 * \|H_r t - a\|_2^2 + s.t. + H_c t = b + t \geq 0 + where : + - c is a (dim_a * dim_b, ) metric cost vector (flattened version of C) + - :math:`\gamma = 1/\lambda` is the l2-regularization coefficient + - H_r is a (dim_a, dim_a * dim_b) metric matrix, + which computes the sum along the rows of transport plan T + - H_c is a (dim_b, dim_a * dim_b) metric matrix, + which computes the sum along the columns of transport plan T + - t is a (dim_a * dim_b, ) metric vector (flattened version of 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 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 T + c : np.ndarray (ns * nt, ) + Flattened array of 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 + + +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 + will be added in 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{h_i^T(H_A \phi - y)}{h_i^T H_A \delta - c_i} + where : + - A is the current active set + - h_i is the ith column of auxiliary matrix H + - H_A is the sub-matrix constructed by the columns of H + whose indices belong to the active set A + - c_i is the ith element of cost vector c + - y is the concatenation of source and target distribution + - :math:`\phi` is the intercept of the solutions in current iteration + - :math:`\delta` is the slope of the solutions in current iteration + Parameters + ---------- + phi : np.ndarray (|A|, ) + Intercept of the solutions in current iteration (t is piecewise linear) + delta : np.ndarray (|A|, ) + Slope of the solutions in current iteration (t is piecewise linear) + HtH : np.ndarray (dim_a * dim_b, dim_a * dim_b) + Matrix product of H^T H + Hty : np.ndarray (dim_a + dim_b, ) + Matrix product of H^T y + c: np.ndarray (dim_a * dim_b, ) + Flattened array of cost matrix 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 + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + 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) + + +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{h_{r i}^T(H_{r A} \phi - a) + h_{c i}^T + \phi_u}{h_{r i}^T H_{r A} \delta + h_{c i} \delta_u - c_i} + where : + - A is the current active set + - h_{r i} is the ith column of the matrix H_r + - h_{c i} is the ith column of the matrix H_c + - H_{r A} is the sub-matrix constructed by the columns of H_r + whose indices belong to the active set A + - c_i is the ith element of cost vector c + - y is the concatenation of source and target distribution + - :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 in current + iteration + - :math:`\delta_u` is the slope of Lagrange parameter in current iteration + Parameters + ---------- + phi : np.ndarray (|A|, ) + Intercept of the solutions in current iteration (t is piecewise linear) + delta : np.ndarray (|A|, ) + Slope of the solutions in current iteration (t is piecewise linear) + phi_u : np.ndarray (dim_b, ) + Intercept of the Lagrange parameter in current iteration (also linear) + delta_u : np.ndarray (dim_b, ) + Slope of the Lagrange parameter in current iteration (also linear) + HrHr : np.ndarray (dim_a * dim_b, dim_a * dim_b) + Matrix product of H_r^T H_r + Hc : np.ndarray (dim_b, dim_a * dim_b) + Matrix that computes the sum along the columns of transport plan T + Hra : np.ndarray (dim_a * dim_b, ) + Matrix product of H_r^T a + c: np.ndarray (dim_a * dim_b, ) + Flattened array of cost matrix 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 + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + + 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) + + +def compute_next_removal(phi, delta, current_gamma): + r""" This function computes the next value of gamma if a variable + is removed in next iteration of regularization path + + We look for the largest value of gamma such that + an element of current solution vanishes + .. math:: + \max_{j \in A} \frac{\phi_j}{\delta_j} + where : + - A is the current active set + - phi_j is the jth element of the intercept of current solution + - delta_j is the jth elemnt of the slope of current solution + Parameters + ---------- + phi : np.ndarray (|A|, ) + Intercept of the solutions in current iteration (t is piecewise linear) + delta : np.ndarray (|A|, ) + Slope of the solutions in current iteration (t is piecewise linear) + current_gamma : float + Value of regularization coefficient at the start of current iteration + Returns + ------- + next_removal_gamma : float + Value of gamma if a variable is removed in next iteration + next_removal_index : int + Index of the variable to remove in next iteration + References + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + 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) + + +def complement_schur(M_current, b, d, id_pop): + r""" This function computes the inverse of matrix in regularization path + using Schur complement + + Two cases may arise: Firstly one variable is added to the active set + .. math:: + M_{k+1}^{-1} = + \begin{bmatrix} + M_{k}^{-1} + s^{-1} M_{k}^{-1} b b^T M_{k}^{-1} & -s^{-1} \\ + - s^{-1} b^T M_{k}^{-1} & s^{-1} + \end{bmatrix} + where : + - :math:`M_k^{-1}` is the inverse of matrix in previous iteration and + :math:`M_k` is the upper left block matrix in Schur formulation + - b is the upper right block matrix in Schur formulation. In our case, + b is reduced to a column vector and b^T is the lower left block matrix + - s is the Schur complement, given by + :math:`s = d - b^T M_{k}^{-1} b` in our case + + Secondly, one variable is removed from the active set + .. math:: + M_{k+1}^{-1} = M^{-1}_{A_k \backslash q} - + \frac{r_{-q,q} r^{T}_{-q,q}}{r_{q,q}} + where : + - q is the index of column and row to delete + - :math:`M^{-1}_{A_k \backslash q}` is the previous inverse matrix + without qth column and qth row + - r_{-q,q} is the qth column of :math:`M^{-1}_{k}` without the qth element + - r_{q, q} is the element of qth column and qth row in :math:`M^{-1}_{k}` + Parameters + ---------- + M_current : np.ndarray (|A|-1, |A|-1) + Inverse matrix in previous iteration + b : np.ndarray (|A|-1, ) + Upper right matrix in Schur complement, a column vector in our case + d : float + Lower right matrix in Schur complement, a scalar in our case + id_pop + Index of the variable to be removed, equal to -1 + if none of the variables is deleted in current iteration + Returns + ------- + M : np.ndarray (|A|, |A|) + Inverse matrix needed in current iteration + References + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + 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 + + +def construct_augmented_H(active_index, m, Hc, HrHr): + r""" This function construct an augmented matrix for the first iteration of + semi-relaxed regularization path + + .. math:: + Augmented_H = + \begin{bmatrix} + 0 & H_{c A} \\ + H_{c A}^T & H_{r A}^T H_{r A} + \end{bmatrix} + where : + - H_{r A} is the sub-matrix constructed by the columns of H_r + whose indices belong to the active set A + - H_{c A} is the sub-matrix constructed by the columns of H_c + whose indices belong to the active set A + Parameters + ---------- + active_index : list + Indices of 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 transport plan T + HrHr : np.ndarray (dim_a * dim_b, dim_a * dim_b) + Matrix product of H_r^T H_r + Returns + ------- + H_augmented : np.ndarray (dim_b + |A|, dim_b + |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 + + +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 c^T t + 0.5 * \|H t - y\|_2^2 + s.t. + t \geq 0 + where : + - c is a (dim_a * dim_b, ) metric cost vector (flattened version of C) + - :math:`\gamma = 1/\lambda` is the l2-regularization coefficient + - y is the concatenation of vectors a and b, defined as y^T = [a^T b^T] + - H is a (dim_a + dim_b, dim_a * dim_b) metric matrix, + see [Chapel et al., 2021] for the design of H. The matrix product Ht + computes both the source marginal and the target marginal. + - t is a (dim_a * dim_b, ) metric vector (flattened version of 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 + l2-regularization coefficient + itmax: int + Maximum number of iteration + Returns + ------- + t : np.ndarray (dim_a*dim_b, ) + Flattened vector of optimal transport matrix + t_list : list + List of solutions in regularization path + gamma_list : list + List of regularization coefficient in 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 + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + + 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., 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 + + +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 c^T t + 0.5 * \|H_r t - a\|_2^2 + s.t. + H_c t = b + t \geq 0 + where : + - c is a (dim_a * dim_b, ) metric cost vector (flattened version of C) + - :math:`\gamma = 1/\lambda` is the l2-regularization coefficient + - H_r is a (dim_a, dim_a * dim_b) metric matrix, + which computes the sum along the rows of transport plan T + - H_c is a (dim_b, dim_a * dim_b) metric matrix, + which computes the sum along the columns of transport plan T + - t is a (dim_a * dim_b, ) metric vector (flattened version of 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 optimal transport matrix + t_list : list + List of solutions in regularization path + gamma_list : list + List of regularization coefficient in 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.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 + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + + 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., 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 + + +def regularization_path(a: np.array, b: np.array, C: np.array, reg=1e-4, + semi_relaxed=False, itmax=50000): + r"""This function combines both the semi-relaxed and the fully-relaxed + regularization paths of l2-UOT problem + + 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 optimal transport matrix + t_list : list + List of solutions in regularization path + gamma_list : list + List of regularization coefficient in regularization path + References + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + 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 + + +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 by the piecewise linearity of the path + + .. math:: + t(\gamma) = \phi(\gamma) - \gamma \delta(\gamma) + where : + - :math:`\gamma` is the regularization coefficient + - :math:`\phi(\gamma)` is the corresponding intercept + - :math:`\delta(\gamma)` is the corresponding slope + - t is a (dim_a * dim_b, ) vector (flattened version of transport matrix) + Parameters + ---------- + gamma : float + Regularization coefficient + gamma_list : list + List of regularization coefficients in regularization path + Pi_list : list + List of solutions in regularization path + Returns + ------- + t : np.ndarray (dim_a*dim_b, ) + Transport vector 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 + ---------- + [Chapel et al., 2021]: + Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021). + Unbalanced optimal transport through non-negative penalized + linear regression. + """ + + 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 |