enable HSGPs for the exponential kernel #234

DESCRIPTION

@@ -2,8 +2,8 @@ Package: brms
 Encoding: UTF-8
 Type: Package
 Title: Bayesian Regression Models using 'Stan'
-Version: 2.22.0
-Date: 2024-09-20
+Version: 2.22.1
+Date: 2024-09-24
 Authors@R:
     c(person("Paul-Christian", "Bürkner", email = "[email protected]",
              role = c("aut", "cre")),

R/formula-gp.R

@@ -16,7 +16,7 @@
 #' @param cov Name of the covariance kernel. Currently supported are
 #'   \code{"exp_quad"} (exponentiated-quadratic kernel; default),
 #'   \code{"matern32"} (Matern 3/2 kernel), \code{"matern52"} (Matern 5/2 kernel),
-#'   and \code{"exponential"} (exponential kernel).
+#'   and \code{"exponential"} (exponential kernel; alias: \code{"matern12"}).
 #' @param iso A flag to indicate whether an isotropic (\code{TRUE}; the
 #'   default) or a non-isotropic GP should be used.
 #'   In the former case, the same amount of smoothing is applied to all
@@ -107,11 +107,10 @@
 #' @export
 gp <- function(..., by = NA, k = NA, cov = "exp_quad", iso = TRUE,
                gr = TRUE, cmc = TRUE, scale = TRUE, c = 5/4) {
-  cov_choices <- c("exp_quad", "matern52", "matern32", "exponential")
-  cov <- match.arg(cov, choices = cov_choices)
   call <- match.call()
   label <- deparse0(call)
   vars <- as.list(substitute(list(...)))[-1]
+  cov <- validate_gp_cov(cov, k = k)
   by <- deparse0(substitute(by))
   cmc <- as_one_logical(cmc)
   if (is.null(call[["gr"]]) && require_old_default("2.12.8")) {
@@ -126,10 +125,6 @@ gp <- function(..., by = NA, k = NA, cov = "exp_quad", iso = TRUE,
     iso <- TRUE
   }
   if (!isNA(k)) {
-    supported_hsgp_covs <- c("exp_quad", "matern52", "matern32")
-    if (!cov %in% supported_hsgp_covs) {
-      stop2("HSGPs with covariance kernel '", cov, "' are not yet supported.")
-    }
     k <- as.integer(as_one_numeric(k))
     if (k < 1L) {
       stop2("'k' must be positive.")
@@ -322,6 +317,34 @@ spd_gp_exp_quad <- function(x, sdgp = 1, lscale = 1) {
   out
 }
 
+# spectral density function of the exponential kernel
+# vectorized over parameter values
+spd_gp_exponential <- function(x, sdgp = 1, lscale = 1) {
+  NB <- NROW(x)
+  D <- NCOL(x)
+  Dls <- NCOL(lscale)
+  constant = square(sdgp) *
+    (2^D * pi^(D / 2) * gamma((D + 1) / 2)) / sqrt(pi)
+  expo = -(D + 1) / 2
+  lscale2 <- lscale^2
+  out <- matrix(nrow = length(sdgp), ncol = NB)
+  if (Dls == 1L) {
+    # one dimensional or isotropic GP
+    constant <- constant * lscale^D
+    for (m in seq_len(NB)) {
+      out[, m] <- constant * (1 + lscale2 * sum(x[m, ]^2))^expo;
+    }
+  } else {
+    # multi-dimensional non-isotropic GP
+    constant <- constant * matrixStats::rowProds(lscale)
+    for (m in seq_len(NB)) {
+      x2 <- data2draws(x[m, ]^2, dim = dim(lscale))
+      out[, m] <- constant * (1 + rowSums(lscale2 * x2))^expo
+    }
+  }
+  out
+}
+
 # spectral density function of the Matern 3/2 kernel
 # vectorized over parameter values
 spd_gp_matern32 <- function(x, sdgp = 1, lscale = 1) {
@@ -406,6 +429,28 @@ choose_L <- function(x, c) {
   c * range
 }
 
+# validate the 'cov' argument of 'gp' terms
+validate_gp_cov <- function(cov, k = NA) {
+  cov <- as_one_character(cov)
+  if (cov == "matern12") {
+    # matern12 and exponential refer to the same kernel
+    cov <- "exponential"
+  }
+  cov_choices <- c("exp_quad", "matern52", "matern32", "exponential")
+  if (!cov %in% cov_choices) {
+    stop2("'", cov, "' is not a valid GP covariance kernel. Valid kernels are: ",
+          collapse_comma(cov_choices))
+  }
+  if (!isNA(k)) {
+    # currently all kernels support HSGPs but this may change in the future
+    hs_cov_choices <- c("exp_quad", "matern52", "matern32", "exponential")
+    if (!cov %in% hs_cov_choices) {
+      stop2("HSGPs with covariance kernel '", cov, "' are not yet supported.")
+    }
+  }
+  cov
+}
+
 # try to evaluate a GP term and
 # return an informative error message if it fails
 try_nug <- function(expr, nug) {

inst/chunks/fun_gp_exponential.stan

@@ -1,4 +1,5 @@
   /* compute a latent Gaussian process with exponential kernel
+   * also known as the Matern 1/2 kernel
    * Args:
    *   x: array of continuous predictor values
    *   sdgp: marginal SD parameter

inst/chunks/fun_spd_gp_exponential.stan

@@ -0,0 +1,34 @@
+  /* Spectral density function of a Gaussian process with exponential kernel
+   * also known as the Matern 1/2 kernel
+   * Args:
+   *   x: array of numeric values of dimension NB x D
+   *   sdgp: marginal SD parameter
+   *   lscale: vector of length-scale parameters
+   * Returns:
+   *   numeric vector of length NB of the SPD evaluated at 'x'
+   */
+  vector spd_gp_exponential(data array[] vector x, real sdgp, vector lscale) {
+    int NB = dims(x)[1];
+    int D = dims(x)[2];
+    int Dls = rows(lscale);
+    real constant = square(sdgp) *
+      (2^D * pi()^(D / 2.0) * tgamma((D + 1.0) / 2)) / sqrt(pi());
+    real expo = -(D + 1.0) / 2;
+    vector[NB] out;
+    if (Dls == 1) {
+      // one dimensional or isotropic GP
+      real lscale2 = square(lscale[1]);
+      constant = constant * lscale[1]^D;
+      for (m in 1:NB) {
+        out[m] = constant * (1 + lscale2 * dot_self(x[m]))^expo;
+      }
+    } else {
+      // multi-dimensional non-isotropic GP
+      vector[Dls] lscale2 = square(lscale);
+      constant = constant * prod(lscale);
+      for (m in 1:NB) {
+        out[m] = constant * (1 + dot_product(lscale2, square(x[m])))^expo;
+      }
+    }
+    return out;
+  }