(Sliced) Wasserstein - Gaussian process regression workflow • slicer

Goal

This vignette illustrates the main use of the slicer package by means of a toy example. Empirical distributions are generated from bivariate Gaussian mixtures with four isotropic components laid out symmetrically around the origin. These distributions differ from one another in that they are rotated counterclockwise by an angle $\omega \in [0, \pi/2)$ . The goal of the analysis is to retrieve the specific angles by which the distributions have been rotated.

Functions

Let’s load slicer.

library(slicer)

To generate the input distributions, we define generate_data():

generate_data <- function(n = 200, a = 5, angle = 0, Sigma = diag(1, 2)) {
  mus <- cbind(c(a, 0), c(0, a), c(-a, 0), c(0, -a))
  R <- rbind(
    c(cos(angle), -sin(angle)),
    c(sin(angle), cos(angle))
  )
  cluster <- sample(1:4, size = n, replace = TRUE)
  d <- matrix(0, nrow = n, ncol = 2)
  for (i in 1:4) {
    d[which(cluster == i), ] <- MASS::mvrnorm(
      sum(cluster == i), mu = mus[, i], Sigma = Sigma
    )
  }
  d %*% t(R)
}

The plots below show two empirical distributions, once using a rotation angle of $\omega = \pi/16$ and once using one of $\omega = 3\pi/8$ .

set.seed(2026) # for reproducibility

op <- par(no.readonly = TRUE)
par(mfrow = c(1, 2))
generate_data(angle = pi/16) |> 
  plot(xlab = "first dimension", ylab = "second dimension", asp = 1)
generate_data(angle = 3*pi/8) |> 
  plot(xlab = "first dimension", ylab = "second dimension", asp = 1)

par(op)

Data generation

We generate 60 matrices (40 for training, 20 for testing) containing empirical distributions with the rotation angles sampled uniformly from $[0, \pi/4)$ . Each empirical distribution is based on 200 data points.

sample_size <- 200
N_train <- 40
N_test  <- 20
angle_range <- c(0, pi/4)

angles <- runif(N_train + N_test, angle_range[1], angle_range[2])

distributions <- vector("list", N_train + N_test)
for (i in seq_len(N_train + N_test)) {
  distributions[[i]] <- generate_data(angle = angles[i])
}

Distances

We can use compute_all_distances() to compute pairwise distances between the input distributions. If we supply the optional vector test_idx with the test set indices, the pairwise distance computations between test objects (which are not needed for obtaining predictions) are skipped. To compute sliced Wasserstein distances, we need to first generate $\boldsymbol{\theta_1}, \dots, \boldsymbol{\theta_L}$ sampled uniformly at random from the unit sphere; generate_directions() takes care of this. We’ll use $L = 25$ and set $d = 2$ since we’re working in two dimensions. If we’re only interested in the estimated sliced Wasserstein distances, we can set keep_projections to FALSE.

thetas <- generate_directions(L = 25, d = 2)
sw_distances <- compute_all_distances(distributions, thetas, verbose = FALSE,
    keep_projections = FALSE, test_idx = N_train + seq_len(N_test))

The $60 \times 60$ matrix sw_distances contains the squared (estimated) sliced Wasserstein distances among the forty training objects and between the training objects and the twenty test objects. The pairwise distances among the test objects are set to NA:

sw_distances # output not shown

We can also compute (normal) Wasserstein distances between the distributions when they are projected along certain dimensions. For instances, to compute the pairwise Wasserstein distances along the first margin and along the second margin, we set the projection directions to $\boldsymbol{e_1} = (1, 0)^{\top}$ and $\boldsymbol{e_2} = (0, 1)^{\top}$ . We also set keep_projections = TRUE, which will cause the function to output a list of two matrices with squared Wasserstein distances: one for each margin.

marginal_distances <- compute_all_distances(distributions, diag(1, 2), 
    verbose = FALSE, keep_projections = TRUE, 
    test_idx = N_train + seq_len(N_test))
str(marginal_distances)
#> List of 2
#>  $ : num [1:60, 1:60] 0 1.274 0.931 0.548 0.649 ...
#>  $ : num [1:60, 1:60] 0 2.177 0.121 0.445 0.599 ...

Gaussian process models with tuned hyperparameters

The function fit_gpr() takes a single matrix with squared pairwise distances and uses it as input to a Gaussian process regression model with (by default) a Gaussian RBF kernel: $k_{\textrm{RBF}}(x, x') = s^2\exp\left(-\frac{d^2(x, x')}{2t^2}\right)$ , where $s^2 > 0$ is a scaling factor, $t > 0$ is the length-scale, and $d(x, x')$ is the distance between the inputs $x$ and $x'$ .

The model’s and the kernel’s hyperparameters are tuned using the training data by minimising the negative marginal log-likelihood.

sw_fit <- fit_gpr(sw_distances,
  training_idx = seq_len(N_train),
  test_idx = N_train + seq_len(N_test),
  y_train = angles[seq_len(N_train)], 
  y_test = angles[N_train + seq_len(N_test)],
  verbose = TRUE)
#> Hyperparameter search 1 of 10.
#> Optimum set at -58.8175077.
#> Hyperparameter search 2 of 10.
#> Hyperparameter search 3 of 10.
#> Hyperparameter search 4 of 10.
#> Hyperparameter search 5 of 10.
#> Hyperparameter search 6 of 10.
#> Hyperparameter search 7 of 10.
#> Hyperparameter search 8 of 10.
#> Current optimum improved from -58.8175077 to -58.8176617.
#> Hyperparameter search 9 of 10.
#> Hyperparameter search 10 of 10.

The output includes predictions for the test data, the covariannce for the predictions (if the true test outcomes were provided), and the estimated hyperparameters.

str(sw_fit)
#> List of 8
#>  $ kernels         : chr "rbf"
#>  $ test_predictions: num [1:20] 0.241 0.7091 0.2402 0.5074 0.0604 ...
#>  $ test_variance   : num [1:20, 1:20] 0.00134 NA NA NA NA ...
#>  $ RMSE            : num 0.038
#>  $ length_scale    : num 1.57
#>  $ scaling_factor  : num 0.0758
#>  $ noise_variance  : num 1.57e-07
#>  $ nll             : num -58.8
plot(angles[N_train + seq_len(N_test)], sw_fit$test_predictions,
     xlab = "true test outcomes", ylab = "predicted test outcomes", asp = 1)

Since we didn’t compute the distances between the test inputs, only the predictions variances are reported (the main diagonal of test_variance):

diag(sw_fit$test_variance)
#>  [1] 0.0013366282 0.0021746730 0.0009410588 0.0024842495 0.0009932353
#>  [6] 0.0012133334 0.0026980876 0.0017291003 0.0009257171 0.0010941720
#> [11] 0.0014113678 0.0014527097 0.0008645677 0.0009060975 0.0030615151
#> [16] 0.0008919646 0.0017044887 0.0039857607 0.0025413811 0.0023768280

Had we also computed the distances between the test inputs, test_variance would have also included the covariances.

The function fit_gpr() can also be used when multiple matrices with squared distances are provided. Now, estimated length-scale and (kernel) variance hyperparameters are provided for the Gaussian RBF corresponding to each squared distance matrix.

marginal_fit <- fit_gpr(marginal_distances,
  training_idx = seq_len(N_train),
  test_idx = N_train + seq_len(N_test),
  y_train = angles[seq_len(N_train)], 
  y_test = angles[N_train + seq_len(N_test)],
  verbose = TRUE)
#> Hyperparameter search 1 of 10.
#> Optimum set at -48.3151392.
#> Hyperparameter search 2 of 10.
#> Current optimum improved from -48.3151392 to -61.5134707.
#> Hyperparameter search 3 of 10.
#> Hyperparameter search 4 of 10.
#> Hyperparameter search 5 of 10.
#> Hyperparameter search 6 of 10.
#> Hyperparameter search 7 of 10.
#> Hyperparameter search 8 of 10.
#> Current optimum improved from -61.5134707 to -61.5134707.
#> Hyperparameter search 9 of 10.
#> Hyperparameter search 10 of 10.
str(marginal_fit)
#> List of 8
#>  $ kernels         : chr [1:2] "rbf" "rbf"
#>  $ test_predictions: num [1:20] 0.2395 0.6678 0.2429 0.5619 0.0919 ...
#>  $ test_variance   : num [1:20, 1:20] 0.000587 NA NA NA NA ...
#>  $ RMSE            : num 0.0522
#>  $ length_scale    : num [1:2] 3.9 6.09
#>  $ scaling_factor  : num [1:2] 0.347 0.146
#>  $ noise_variance  : num 0.000448
#>  $ nll             : num -61.5
plot(angles[N_train + seq_len(N_test)], marginal_fit$test_predictions,
     xlab = "true test outcomes", ylab = "predicted test outcomes", asp = 1)

We can combine the marginal and sliced Wasserstein distances into a list with three distance matrices, too:

total_fit <- fit_gpr(list(sw_distances, marginal_distances[[1]], marginal_distances[[2]]),
  training_idx = seq_len(N_train),
  test_idx = N_train + seq_len(N_test),
  y_train = angles[seq_len(N_train)], 
  y_test = angles[N_train + seq_len(N_test)],
  verbose = TRUE)
#> Hyperparameter search 1 of 10.
#> Optimum set at -58.8320454.
#> Hyperparameter search 2 of 10.
#> Current optimum improved from -58.8320454 to -61.5129334.
#> Hyperparameter search 3 of 10.
#> Current optimum improved from -61.5129334 to -63.080251.
#> Hyperparameter search 4 of 10.
#> Hyperparameter search 5 of 10.
#> Hyperparameter search 6 of 10.
#> Hyperparameter search 7 of 10.
#> Hyperparameter search 8 of 10.
#> Hyperparameter search 9 of 10.
#> Hyperparameter search 10 of 10.
str(total_fit)
#> List of 8
#>  $ kernels         : chr [1:3] "rbf" "rbf" "rbf"
#>  $ test_predictions: num [1:20] 0.247 0.69 0.247 0.549 0.085 ...
#>  $ test_variance   : num [1:20, 1:20] 0.00081 NA NA NA NA ...
#>  $ RMSE            : num 0.047
#>  $ length_scale    : num [1:3] 0.421 3.833 5.482
#>  $ scaling_factor  : num [1:3] 0.00118 0.24703 0.12446
#>  $ noise_variance  : num 3.32e-08
#>  $ nll             : num -63.1
plot(angles[N_train + seq_len(N_test)], total_fit$test_predictions,
     xlab = "true test outcomes", ylab = "predicted test outcomes", asp = 1)

Parallel processing can be enabled using the cores parameter:

total_fit <- fit_gpr(list(sw_distances, marginal_distances[[1]], marginal_distances[[2]]),
  training_idx = seq_len(N_train),
  test_idx = N_train + seq_len(N_test),
  y_train = angles[seq_len(N_train)], 
  y_test = angles[N_train + seq_len(N_test)],
  runs = 50L, cores = 2L)
str(total_fit)
#> List of 8
#>  $ kernels         : chr [1:3] "rbf" "rbf" "rbf"
#>  $ test_predictions: num [1:20] 0.247 0.69 0.247 0.549 0.085 ...
#>  $ test_variance   : num [1:20, 1:20] 0.00081 NA NA NA NA ...
#>  $ RMSE            : num 0.047
#>  $ length_scale    : num [1:3] 0.421 3.837 5.494
#>  $ scaling_factor  : num [1:3] 0.00117 0.24769 0.12513
#>  $ noise_variance  : num 1.09e-07
#>  $ nll             : num -63.1

Instead of the Gaussian RBF kernel, the Matérn kernels with smoothness parameter $\nu = 0.5, 1.5, 2.5$ can be used. We can use a different kernel for each distance matrix; the total kernel used is then the conical combination of these kernels. For instance, the following fit uses the Matérn kernel with $\nu = 2.5$ for the sliced Wasserstein distances, the Matérn kernel with $\nu = 0.5$ for the Wasserstein distances along the first axis, and the Matérn kernel with $\nu = 1.5$ for the Wasserstein distances alogn the second axis.

total_fit_matern <- fit_gpr(list(sw_distances, marginal_distances[[1]], marginal_distances[[2]]),
  training_idx = seq_len(N_train),
  test_idx = N_train + seq_len(N_test),
  y_train = angles[seq_len(N_train)], 
  y_test = angles[N_train + seq_len(N_test)],
  kernels = c("matern25", "matern05", "matern15"),
  runs = 50L, cores = 2L)
str(total_fit_matern)
#> List of 8
#>  $ kernels         : chr [1:3] "matern25" "matern05" "matern15"
#>  $ test_predictions: num [1:20] 0.242 0.709 0.24 0.507 0.061 ...
#>  $ test_variance   : num [1:20, 1:20] 0.00133 NA NA NA NA ...
#>  $ RMSE            : num 0.0377
#>  $ length_scale    : num [1:3] 2.71 1.69e+04 2.03e+07
#>  $ scaling_factor  : num [1:3] 1.34e-01 7.23e-09 4.17e-10
#>  $ noise_variance  : num 4.96e-09
#>  $ nll             : num -58.6

Finally, a GPR’s negative log predictive density (NLPD) for some test data can be obtained like so:

nlpd_gpr(total_fit_matern, angles[N_train + seq_len(N_test)])
#> Warning in nlpd_gpr(total_fit_matern, angles[N_train + seq_len(N_test)]): The fit's test variance contains NAs, likely because the distances among the test inputs weren't provided.
#> Replacing NAs by 0.
#> [1] -1.955064
#> attr(,"jitter_used")
#> [1] 0

The warning message occurs because we didn’t compute the distances among the test inputs themselves. As a result, the kernel matrices don’t specify the covariance among the test inputs, and the NLPD assumes that they are zero. If you remove the test_idx parameter setting when running compute_all_distances() and refit the models, this warning will disappear.