Source code for ot.lp

# -*- coding: utf-8 -*-
"""
Solvers for the original linear program OT problem
"""

# Author: Remi Flamary <remi.flamary@unice.fr>
#
# License: MIT License

import multiprocessing

import numpy as np

from .import cvx

# import compiled emd
from .emd_wrap import emd_c, check_result
from ..utils import parmap
from .cvx import barycenter
from ..utils import dist

__all__=['emd', 'emd2', 'barycenter', 'free_support_barycenter', 'cvx']


[docs]def emd(a, b, M, numItermax=100000, log=False): """Solves the Earth Movers distance problem and returns the OT matrix .. math:: \gamma = arg\min_\gamma <\gamma,M>_F s.t. \gamma 1 = a \gamma^T 1= b \gamma\geq 0 where : - M is the metric cost matrix - a and b are the sample weights Uses the algorithm proposed in [1]_ Parameters ---------- a : (ns,) ndarray, float64 Source histogram (uniform weigth if empty list) b : (nt,) ndarray, float64 Target histogram (uniform weigth if empty list) M : (ns,nt) ndarray, float64 loss matrix numItermax : int, optional (default=100000) The maximum number of iterations before stopping the optimization algorithm if it has not converged. log: boolean, optional (default=False) If True, returns a dictionary containing the cost and dual variables. Otherwise returns only the optimal transportation matrix. Returns ------- gamma: (ns x nt) ndarray Optimal transportation matrix for the given parameters log: dict If input log is true, a dictionary containing the cost and dual variables and exit status Examples -------- Simple example with obvious solution. The function emd accepts lists and perform automatic conversion to numpy arrays >>> import ot >>> a=[.5,.5] >>> b=[.5,.5] >>> M=[[0.,1.],[1.,0.]] >>> ot.emd(a,b,M) array([[ 0.5, 0. ], [ 0. , 0.5]]) References ---------- .. [1] Bonneel, N., Van De Panne, M., Paris, S., & Heidrich, W. (2011, December). Displacement interpolation using Lagrangian mass transport. In ACM Transactions on Graphics (TOG) (Vol. 30, No. 6, p. 158). ACM. See Also -------- ot.bregman.sinkhorn : Entropic regularized OT ot.optim.cg : General regularized OT""" a = np.asarray(a, dtype=np.float64) b = np.asarray(b, dtype=np.float64) M = np.asarray(M, dtype=np.float64) # if empty array given then use unifor distributions if len(a) == 0: a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0] if len(b) == 0: b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1] G, cost, u, v, result_code = emd_c(a, b, M, numItermax) result_code_string = check_result(result_code) if log: log = {} log['cost'] = cost log['u'] = u log['v'] = v log['warning'] = result_code_string log['result_code'] = result_code return G, log return G
[docs]def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000, log=False, return_matrix=False): """Solves the Earth Movers distance problem and returns the loss .. math:: \gamma = arg\min_\gamma <\gamma,M>_F s.t. \gamma 1 = a \gamma^T 1= b \gamma\geq 0 where : - M is the metric cost matrix - a and b are the sample weights Uses the algorithm proposed in [1]_ Parameters ---------- a : (ns,) ndarray, float64 Source histogram (uniform weigth if empty list) b : (nt,) ndarray, float64 Target histogram (uniform weigth if empty list) M : (ns,nt) ndarray, float64 loss matrix numItermax : int, optional (default=100000) The maximum number of iterations before stopping the optimization algorithm if it has not converged. log: boolean, optional (default=False) If True, returns a dictionary containing the cost and dual variables. Otherwise returns only the optimal transportation cost. return_matrix: boolean, optional (default=False) If True, returns the optimal transportation matrix in the log. Returns ------- gamma: (ns x nt) ndarray Optimal transportation matrix for the given parameters log: dict If input log is true, a dictionary containing the cost and dual variables and exit status Examples -------- Simple example with obvious solution. The function emd accepts lists and perform automatic conversion to numpy arrays >>> import ot >>> a=[.5,.5] >>> b=[.5,.5] >>> M=[[0.,1.],[1.,0.]] >>> ot.emd2(a,b,M) 0.0 References ---------- .. [1] Bonneel, N., Van De Panne, M., Paris, S., & Heidrich, W. (2011, December). Displacement interpolation using Lagrangian mass transport. In ACM Transactions on Graphics (TOG) (Vol. 30, No. 6, p. 158). ACM. See Also -------- ot.bregman.sinkhorn : Entropic regularized OT ot.optim.cg : General regularized OT""" a = np.asarray(a, dtype=np.float64) b = np.asarray(b, dtype=np.float64) M = np.asarray(M, dtype=np.float64) # if empty array given then use unifor distributions if len(a) == 0: a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0] if len(b) == 0: b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1] if log or return_matrix: def f(b): G, cost, u, v, resultCode = emd_c(a, b, M, numItermax) result_code_string = check_result(resultCode) log = {} if return_matrix: log['G'] = G log['u'] = u log['v'] = v log['warning'] = result_code_string log['result_code'] = resultCode return [cost, log] else: def f(b): G, cost, u, v, result_code = emd_c(a, b, M, numItermax) check_result(result_code) return cost if len(b.shape) == 1: return f(b) nb = b.shape[1] res = parmap(f, [b[:, i] for i in range(nb)], processes) return res
[docs]def free_support_barycenter(measures_locations, measures_weights, X_init, b=None, weights=None, numItermax=100, stopThr=1e-7, verbose=False, log=None): """ 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) The function solves the Wasserstein barycenter problem when the barycenter measure is constrained to be supported on k atoms. This problem is considered in [1] (Algorithm 2). There are two differences with the following codes: - we do not optimize over the weights - 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. Parameters ---------- measures_locations : list of (k_i,d) np.ndarray 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) measures_weights : list of (k_i,) np.ndarray Numpy arrays where each numpy array has k_i non-negatives values summing to one representing the weights of each discrete input measure X_init : (k,d) np.ndarray Initialization of the support locations (on k atoms) of the barycenter b : (k,) np.ndarray Initialization of the weights of the barycenter (non-negatives, sum to 1) weights : (k,) np.ndarray Initialization of the coefficients of the barycenter (non-negatives, sum to 1) numItermax : int, optional Max number of iterations stopThr : float, optional Stop threshol on error (>0) verbose : bool, optional Print information along iterations log : bool, optional record log if True Returns ------- X : (k,d) np.ndarray Support locations (on k atoms) of the barycenter References ---------- .. [1] Cuturi, Marco, and Arnaud Doucet. "Fast computation of Wasserstein barycenters." International Conference on Machine Learning. 2014. .. [2] Álvarez-Esteban, Pedro C., et al. "A fixed-point approach to barycenters in Wasserstein space." Journal of Mathematical Analysis and Applications 441.2 (2016): 744-762. """ iter_count = 0 N = len(measures_locations) k = X_init.shape[0] d = X_init.shape[1] if b is None: b = np.ones((k,))/k if weights is None: weights = np.ones((N,)) / N X = X_init log_dict = {} displacement_square_norms = [] displacement_square_norm = stopThr + 1. while ( displacement_square_norm > stopThr and iter_count < numItermax ): T_sum = np.zeros((k, d)) for (measure_locations_i, measure_weights_i, weight_i) in zip(measures_locations, measures_weights, weights.tolist()): M_i = dist(X, measure_locations_i) T_i = emd(b, measure_weights_i, M_i) T_sum = T_sum + weight_i * np.reshape(1. / b, (-1, 1)) * np.matmul(T_i, measure_locations_i) displacement_square_norm = np.sum(np.square(T_sum-X)) if log: displacement_square_norms.append(displacement_square_norm) X = T_sum if verbose: print('iteration %d, displacement_square_norm=%f\n', iter_count, displacement_square_norm) iter_count += 1 if log: log_dict['displacement_square_norms'] = displacement_square_norms return X, log_dict else: return X