Procrustes-Wasserstein Distance — pwdist • T4transport

Given two empirical measures $$\mu = \sum_{m=1}^M \mu_m \delta_{X_m}\quad\textrm{and}\quad \nu = \sum_{n=1}^N \nu_n \delta_{Y_n}$$ in $\mathbb{R}^P$, the Procrustes-Wasserstein (PW) distance is defined as follows: $$ PW_2^2(\mu, \nu) = \min_{Q\in \mathcal{O}(P)} W_2^2(\mu, Q_\# \nu), $$ where $\mathcal{O}(P)$ is the orthogonal group and $Q_\#\nu$ is the pushforward via $Q$.

Alternating exact transport and orthogonal Procrustes updates give a local comparison; global minimization over alignments is not certified. The algorithm does not translate, center, or rescale the observations. Accepted objective history, initialization, controls, and termination status are returned. The default compares both initial alignments using successful exact solves only. A failed alternative is recorded in initialization$attempts without replacing an available valid start. If both starts fail, no plan or attained objective is returned.

Usage

pwdist(X, Y, wx = NULL, wy = NULL, ...)

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()].

wx

a length-$M$ nonnegative mass vector, normalized internally. If NULL (default), matrix rows have uniform mass. For a measure object this must be NULL and its stored masses are used.

wy

a length-$N$ nonnegative mass vector with the same rules as wx.

...

extra parameters including

maxiter: maximum number of iterations (default: 496).
abstol: absolute tolerance for the change of squared block objective (default: 1e-10).
inner_maxiter: iteration limit for each exact transport solve (default: 100000000; zero means no limit).
init_alignment: optional orthogonal matrix; otherwise compare a weighted index-coupling initialization and identity.

Value

a named list containing

distance: the computed PW distance value.
plan: an $(M\times N)$ nonnegative matrix for the optimal transport plan.
alignment: the attained alignment matrix of size $(P\times P)$ in $\mathcal{O}(P)$.

References

Adamo D, Corneli M, Vuillien M, Vila E (2025). “An in Depth Look at the Procrustes-Wasserstein Distance: Properties and Barycenters.” In Forty-Second International Conference on Machine Learning.

Examples

if (FALSE) { # \dontrun{
#-------------------------------------------------------------------
#                           Description
#
# * class 1 : samples from N((0,0),  diag(c(4,1/4)))
# * class 2 : samples from N((10,0), diag(c(1/4,4)))
# * class 3 : samples from N((10,0), diag(c(1/4,4))) randomly rotated
#
#  We draw 10 empirical measures from each and compare 
#  the regular Wasserstein and PW distance.
#-------------------------------------------------------------------
## GENERATE DATA
set.seed(10)

#  prepare empty lists
inputs = vector("list", length=30)

#  generate
random_rot = qr.Q(qr(matrix(runif(4), ncol=2)))
for (i in 1:10){
  inputs[[i]] = matrix(rnorm(50*2), ncol=2)
}
for (j in 11:20){
  base_draw = matrix(rnorm(50*2), ncol=2)
  base_draw[,1] = base_draw[,1] + 10
  
  inputs[[j]] = base_draw
  inputs[[j+10]] = base_draw%*%random_rot
}

## COMPUTE
#  empty arrays
dist_RW = array(0, c(30, 30))
dist_PW = array(0, c(30, 30))

#  compute pairwise distances
for (i in 1:29){
  for (j in (i+1):30){
  dist_RW[i,j] <- dist_RW[j,i] <- wasserstein(inputs[[i]], inputs[[j]])$distance
  dist_PW[i,j] <- dist_PW[j,i] <- pwdist(inputs[[i]], inputs[[j]])$distance
  }
}

## VISUALIZE
opar <- par(no.readonly=TRUE)
par(mfrow=c(1,2), pty="s")
image(dist_RW, xaxt="n", yaxt="n", main="Regular Wasserstein distance")
image(dist_PW, xaxt="n", yaxt="n", main="PW distance")
par(opar)
} # }