""" Sliced OT Distances """ # Author: Adrien Corenflos # Nicolas Courty # Rémi Flamary # # License: MIT License import numpy as np from .backend import get_backend, NumpyBackend from .utils import list_to_array def get_random_projections(d, n_projections, seed=None, backend=None, type_as=None): r""" Generates n_projections samples from the uniform on the unit sphere of dimension :math:`d-1`: :math:`\mathcal{U}(\mathcal{S}^{d-1})` Parameters ---------- d : int dimension of the space n_projections : int number of samples requested seed: int or RandomState, optional Seed used for numpy random number generator backend: Backend to ue for random generation Returns ------- out: ndarray, shape (d, n_projections) The uniform unit vectors on the sphere Examples -------- >>> n_projections = 100 >>> d = 5 >>> projs = get_random_projections(d, n_projections) >>> np.allclose(np.sum(np.square(projs), 0), 1.) # doctest: +NORMALIZE_WHITESPACE True """ if backend is None: nx = NumpyBackend() else: nx = backend if isinstance(seed, np.random.RandomState) and str(nx) == 'numpy': projections = seed.randn(d, n_projections) else: if seed is not None: nx.seed(seed) projections = nx.randn(d, n_projections, type_as=type_as) projections = projections / nx.sqrt(nx.sum(projections**2, 0, keepdims=True)) return projections def sliced_wasserstein_distance(X_s, X_t, a=None, b=None, n_projections=50, p=2, projections=None, seed=None, log=False): r""" Computes a Monte-Carlo approximation of the p-Sliced Wasserstein distance .. math:: \mathcal{SWD}_p(\mu, \nu) = \underset{\theta \sim \mathcal{U}(\mathbb{S}^{d-1})}{\mathbb{E}}\left(\mathcal{W}_p^p(\theta_\# \mu, \theta_\# \nu)\right)^{\frac{1}{p}} where : - :math:`\theta_\# \mu` stands for the pushforwards of the projection :math:`X \in \mathbb{R}^d \mapsto \langle \theta, X \rangle` Parameters ---------- X_s : ndarray, shape (n_samples_a, dim) samples in the source domain X_t : ndarray, shape (n_samples_b, dim) samples in the target domain a : ndarray, shape (n_samples_a,), optional samples weights in the source domain b : ndarray, shape (n_samples_b,), optional samples weights in the target domain n_projections : int, optional Number of projections used for the Monte-Carlo approximation p: float, optional = Power p used for computing the sliced Wasserstein projections: shape (dim, n_projections), optional Projection matrix (n_projections and seed are not used in this case) seed: int or RandomState or None, optional Seed used for random number generator log: bool, optional if True, sliced_wasserstein_distance returns the projections used and their associated EMD. Returns ------- cost: float Sliced Wasserstein Cost log : dict, optional log dictionary return only if log==True in parameters Examples -------- >>> n_samples_a = 20 >>> reg = 0.1 >>> X = np.random.normal(0., 1., (n_samples_a, 5)) >>> sliced_wasserstein_distance(X, X, seed=0) # doctest: +NORMALIZE_WHITESPACE 0.0 References ---------- .. [31] Bonneel, Nicolas, et al. "Sliced and radon wasserstein barycenters of measures." Journal of Mathematical Imaging and Vision 51.1 (2015): 22-45 """ from .lp import wasserstein_1d X_s, X_t = list_to_array(X_s, X_t) if a is not None and b is not None and projections is None: nx = get_backend(X_s, X_t, a, b) elif a is not None and b is not None and projections is not None: nx = get_backend(X_s, X_t, a, b, projections) elif a is None and b is None and projections is not None: nx = get_backend(X_s, X_t, projections) else: nx = get_backend(X_s, X_t) n = X_s.shape[0] m = X_t.shape[0] if X_s.shape[1] != X_t.shape[1]: raise ValueError( "X_s and X_t must have the same number of dimensions {} and {} respectively given".format(X_s.shape[1], X_t.shape[1])) if a is None: a = nx.full(n, 1 / n, type_as=X_s) if b is None: b = nx.full(m, 1 / m, type_as=X_s) d = X_s.shape[1] if projections is None: projections = get_random_projections(d, n_projections, seed, backend=nx, type_as=X_s) X_s_projections = nx.dot(X_s, projections) X_t_projections = nx.dot(X_t, projections) projected_emd = wasserstein_1d(X_s_projections, X_t_projections, a, b, p=p) res = (nx.sum(projected_emd) / n_projections) ** (1.0 / p) if log: return res, {"projections": projections, "projected_emds": projected_emd} return res def max_sliced_wasserstein_distance(X_s, X_t, a=None, b=None, n_projections=50, p=2, projections=None, seed=None, log=False): r""" Computes a Monte-Carlo approximation of the max p-Sliced Wasserstein distance .. math:: \mathcal{Max-SWD}_p(\mu, \nu) = \underset{\theta _in \mathcal{U}(\mathbb{S}^{d-1})}{\max} [\mathcal{W}_p^p(\theta_\# \mu, \theta_\# \nu)]^{\frac{1}{p}} where : - :math:`\theta_\# \mu` stands for the pushforwars of the projection :math:`\mathbb{R}^d \ni X \mapsto \langle \theta, X \rangle` Parameters ---------- X_s : ndarray, shape (n_samples_a, dim) samples in the source domain X_t : ndarray, shape (n_samples_b, dim) samples in the target domain a : ndarray, shape (n_samples_a,), optional samples weights in the source domain b : ndarray, shape (n_samples_b,), optional samples weights in the target domain n_projections : int, optional Number of projections used for the Monte-Carlo approximation p: float, optional = Power p used for computing the sliced Wasserstein projections: shape (dim, n_projections), optional Projection matrix (n_projections and seed are not used in this case) seed: int or RandomState or None, optional Seed used for random number generator log: bool, optional if True, sliced_wasserstein_distance returns the projections used and their associated EMD. Returns ------- cost: float Sliced Wasserstein Cost log : dict, optional log dictionary return only if log==True in parameters Examples -------- >>> n_samples_a = 20 >>> reg = 0.1 >>> X = np.random.normal(0., 1., (n_samples_a, 5)) >>> sliced_wasserstein_distance(X, X, seed=0) # doctest: +NORMALIZE_WHITESPACE 0.0 References ---------- .. [35] Deshpande, I., Hu, Y. T., Sun, R., Pyrros, A., Siddiqui, N., Koyejo, S., ... & Schwing, A. G. (2019). Max-sliced wasserstein distance and its use for gans. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 10648-10656). """ from .lp import wasserstein_1d X_s, X_t = list_to_array(X_s, X_t) if a is not None and b is not None and projections is None: nx = get_backend(X_s, X_t, a, b) elif a is not None and b is not None and projections is not None: nx = get_backend(X_s, X_t, a, b, projections) elif a is None and b is None and projections is not None: nx = get_backend(X_s, X_t, projections) else: nx = get_backend(X_s, X_t) n = X_s.shape[0] m = X_t.shape[0] if X_s.shape[1] != X_t.shape[1]: raise ValueError( "X_s and X_t must have the same number of dimensions {} and {} respectively given".format(X_s.shape[1], X_t.shape[1])) if a is None: a = nx.full(n, 1 / n, type_as=X_s) if b is None: b = nx.full(m, 1 / m, type_as=X_s) d = X_s.shape[1] if projections is None: projections = get_random_projections(d, n_projections, seed, backend=nx, type_as=X_s) X_s_projections = nx.dot(X_s, projections) X_t_projections = nx.dot(X_t, projections) projected_emd = wasserstein_1d(X_s_projections, X_t_projections, a, b, p=p) res = nx.max(projected_emd) ** (1.0 / p) if log: return res, {"projections": projections, "projected_emds": projected_emd} return res