Sliced Wasserstein Distance — swdist • T4transport

Sliced Wasserstein distance compares univariate projections of probability measures. For the $d$-dimensional probability measures $\mu$ and $\nu$, the SW distance is defined as $$\mathcal{SW}_p (\mu, \nu) = \left( \int_{\mathbb{S}^{d-1}} \mathcal{W}_p^p ( \langle \theta, \mu\rangle, \langle \theta, \nu \rangle) d\lambda (\theta) \right)^{1/p},$$ where $\mathbb{S}^{d-1}$ is the $(d-1)$-dimensional unit hypersphere and $\lambda$ is the uniform distribution on $\mathbb{S}^{d-1}$. Practically, it is approximated by an equal-weight average of powered projected distances, followed by the p-th root. Each projected finite-measure problem is solved by exact monotone transport. A finite collection of directions can define only a pseudometric; no sphere-integration error estimate is returned.

Usage

swdist(X, Y, p = 2, ...)

Arguments

X

an $(M\times P)$ matrix of row observations or a finite measure from [ot_measure()].

Y

an $(N\times P)$ matrix of row observations or a finite measure from [ot_measure()].

p

a finite Wasserstein order at least one (default: 2).

...

extra parameters including

num_proj: number of directions (default: 496, or one in 1D).
directions: optional matrix with one nonzero direction per row and one column per coordinate; rows are normalized to unit length. A numeric vector specifies one direction. If num_proj is also supplied, it must match.
wx,wy: optional nonnegative masses on the rows of X and Y, normalized separately; defaults are uniform. For measure objects, their stored masses are used and the corresponding mass argument must be NULL.

Value

a named list containing

distance: p-th root of the mean powered projected distances.
projdist: a vector of projected univariate distances.
directions: unit directions actually used, one per row, permitting replay.
approximation: exact_1d or finite_direction_average.
status,stop_reason,diagnostics: finite-solver status and direction metadata; success does not certify accuracy of integration over the sphere.

References

Rabin J, Peyré G, Delon J, Bernot M (2012). “Wasserstein Barycenter and Its Application to Texture Mixing.” In Bruckstein AM, ter Haar Romeny BM, Bronstein AM, Bronstein MM (eds.), Scale Space and Variational Methods in Computer Vision, volume 6667, 435–446. Springer Berlin Heidelberg, Berlin, Heidelberg. ISBN 978-3-642-24784-2 978-3-642-24785-9, doi:10.1007/978-3-642-24785-9_37 .

Examples

# \donttest{
#-------------------------------------------------------------------
#  Sliced-Wasserstein Distance between Two Bivariate Normal
#
# * class 1 : samples from Gaussian with mean=(-1, -1)
# * class 2 : samples from Gaussian with mean=(+1, +1)
#-------------------------------------------------------------------
# SMALL EXAMPLE
set.seed(100)
m = 20
n = 30
X = matrix(rnorm(m*2, mean=-1),ncol=2) # m obs. for X
Y = matrix(rnorm(n*2, mean=+1),ncol=2) # n obs. for Y

# COMPUTE THE SLICED-WASSERSTEIN DISTANCE
outsw <- swdist(X, Y, num_proj=100)

# VISUALIZE
# prepare ingredients for plotting
plot_x = seq_along(outsw$projdist)
plot_y = (base::cumsum(outsw$projdist^2)/plot_x)^(1/2)

# draw
opar <- par(no.readonly=TRUE)
plot(plot_x, plot_y, type="b", cex=0.1, lwd=2,
     xlab="number of MC samples", ylab="distance",
     main="Effect of MC Sample Size")
abline(h=outsw$distance, col="red", lwd=2)
legend("bottomright", legend="SW Distance", 
       col="red", lwd=2)

par(opar)
# }