Issue 757: Vectorise GP stan code

R/create.R

@@ -331,6 +331,7 @@ create_backcalc_data <- function(backcalc = backcalc_opts()) {
  )
  return(data)
 }
+
 #' Create Gaussian Process Data
 #'
 #' @description `r lifecycle::badge("stable")`
@@ -394,11 +395,13 @@ create_gp_data <- function(gp = gp_opts(), data) {
     ls_max = data$t - data$seeding_time - data$horizon,
     alpha_sd = gp$alpha_sd,
     gp_type = data.table::fcase(
-      is.infinite(gp$matern_order), 0,
-      gp$matern_order == 1 / 2, 1,
-      gp$matern_order == 3 / 2, 2,
-      default = 3
-    )
+      gp$kernel == "se", 0,
+      gp$kernel == "periodic", 1,
+      gp$kernel == "matern" || gp$kernel == "ou", 2,
+      default = 2
+    ),
+    nu = gp$matern_order,
+    w0 = gp$w0
   )
 
   gp_data <- c(data, gp_data)

R/opts.R

@@ -401,7 +401,7 @@
 #'
 #' @description `r lifecycle::badge("stable")`
 #' Defines a list specifying the structure of the approximate Gaussian
-#'  process. Custom settings can be supplied which override the defaults.
+#' process. Custom settings can be supplied which override the defaults.
 #'
 #' @param ls_mean Numeric, defaults to 21 days. The mean of the lognormal
 #' length scale.
@@ -418,22 +418,20 @@
 #' @param ls_min Numeric, defaults to 0. The minimum value of the length scale.
 #'
 #' @param alpha_sd Numeric, defaults to 0.05. The standard deviation of the
-#'  magnitude parameter of the Gaussian process kernel. Should be approximately
+#' magnitude parameter of the Gaussian process kernel. Should be approximately
 #' the expected standard deviation of the logged Rt.
 #'
 #' @param kernel Character string, the type of kernel required. Currently
-#' supporting the squared exponential kernel ("se", or "matern" with
-#' 'matern_order = Inf'), 3 over 2 oder 5 over 2 Matern kernel ("matern", with
-#' `matern_order = 3/2` (default) or `matern_order = 5/2`, respectively), or
-#' Orstein-Uhlenbeck ("ou", or "matern" with 'matern_order = 1/2'). Defaulting
-#' to the Matérn 3 over 2 kernel for a balance of smoothness and
-#' discontinuities.
+#' supporting the squared exponential kernel ("se"), periodic kernel
+#' ("periodic"), Ornstein-Uhlenbeck kernel ("ou"), and Matern kernel ("matern").
 #'
 #' @param matern_order Numeric, defaults to 3/2. Order of Matérn Kernel to use.
-#' Currently the orders 1/2, 3/2, 5/2 and Inf are supported.
+#' Common choices are 1/2, 3/2, and 5/2. If `kernel` is set
+#' to "ou", `matern_order` will be automatically set to 1/2. Only used if
+#' the kernel is set to "matern".
 #'
-#' @param matern_type Deprated; Numeric, defaults to 3/2. Order of Matérn Kernel
-#'   to use.  Currently the orders 1/2, 3/2, 5/2 and Inf are supported.
+#' @param matern_type Deprecated; Numeric, defaults to 3/2. Order of Matérn Kernel
+#' to use. Currently the orders 1/2, 3/2, 5/2 and Inf are supported.
 #'
 #' @param basis_prop Numeric, proportion of time points to use as basis
 #' functions. Defaults to 0.2. Decreasing this value results in a decrease in
@@ -446,6 +444,9 @@
 #' approximate Gaussian process. See (Riutort-Mayol et al. 2020
 #' <https://arxiv.org/abs/2004.11408>) for advice on updating this default.
 #'
+#' @param w0 Numeric, defaults to 1.0. Fundamental frequency for periodic
+#' kernel. Only used if `kernel` is set to "periodic".
+#'
 #' @importFrom rlang arg_match
 #' @return A `<gp_opts>` object of settings defining the Gaussian process
 #' @export
@@ -462,14 +463,15 @@
                     ls_min = 0,
                     ls_max = 60,
                     alpha_sd = 0.05,
-                    kernel = c("matern", "se", "ou"),
+                    kernel = c("matern", "se", "ou", "periodic"),
                     matern_order = 3 / 2,
-                    matern_type) {
+                    matern_type,
+                    w0 = 1.0) {
   lifecycle::deprecate_warn(
     "1.6.0", "gp_opts(matern_type)", "gp_opts(matern_order)"
   )
   if (!missing(matern_type)) {
-    if (!missing(matern_order) && matern_type == matern_order) {
+    if (!missing(matern_order) && matern_type != matern_order) {
       stop(
         "Incompatible `matern_order` and `matern_type`. ",
         "Use `matern_order` only."
@@ -480,20 +482,15 @@
 
   kernel <- arg_match(kernel)
   if (kernel == "se") {
-    if (!missing(matern_order) && is.finite(matern_order)) {
-      stop("Squared exponential kernel must have matern order unset or `Inf`.")
-    }
     matern_order <- Inf
   } else if (kernel == "ou") {
-    if (!missing(matern_order) && matern_order != 1 / 2) {
-      stop("Ornstein-Uhlenbeck kernel must have matern order unset or `1 / 2`.") ## nolint: nonportable_path_linter
-    }
     matern_order <- 1 / 2
-  } else if (!(is.infinite(matern_order) ||
-                 matern_order %in% c(1 / 2, 3 / 2,  5 / 2))) {
-    stop(
-      "only the Matern kernels of order `1 / 2`, `3 / 2`, `5 / 2` or `Inf` ", ## nolint: nonportable_path_linter
-      "are currently supported"
+  } else if (
+      !(is.infinite(matern_order) || matern_order %in% c(1 / 2, 3 / 2, 5 / 2))
+    ) {
+    warning(
+      "Uncommon Matern kernel order. Common orders are `1 / 2`, `3 / 2`,",
+      " and `5 / 2`"
     )
   }
 
@@ -506,7 +503,8 @@
     ls_max = ls_max,
     alpha_sd = alpha_sd,
     kernel = kernel,
-    matern_order = matern_order
+    matern_order = matern_order,
+    w0 = w0
   )
 
   attr(gp, "class") <- c("gp_opts", class(gp))

inst/stan/data/gaussian_process.stan

@@ -5,6 +5,8 @@
   real<lower=0> ls_min;              // Lower bound for the lengthscale
   real<lower=0> ls_max;              // Upper bound for the lengthscale
   real alpha_sd;                     // standard deviation of the alpha gp kernal parameter
-  int gp_type;                       // type of gp, 0 = squared exponential, 1 = 3/2 matern
+  int gp_type;                       // type of gp, 0 = squared exponential, 1 = periodic, 2 = Matern
+  real nu;                           // smoothness parameter for Matern kernel (used if gp_type = 2)
+  real w0;                           // fundamental frequency for periodic kernel (used if gp_type = 1)
   int stationary;                    // is underlying gaussian process first or second order
   int fixed;                         //  should a gaussian process be used

inst/stan/estimate_infections.stan

@@ -9,7 +9,6 @@ functions {
 #include functions/generated_quantities.stan
 }
 
-
 data {
 #include data/observations.stan
 #include data/delays.stan
@@ -19,15 +18,15 @@ data {
 #include data/observation_model.stan
 }
 
-transformed data{
+transformed data {
   // observations
   int ot = t - seeding_time - horizon;  // observed time
   int ot_h = ot + horizon;  // observed time + forecast horizon
   // gaussian process
   int noise_terms = setup_noise(
     ot_h, t, horizon, estimate_r, stationary, future_fixed, fixed_from
   );
-  matrix[noise_terms, M] PHI = setup_gp(M, L, noise_terms);  // basis function
+  matrix[noise_terms, M] PHI = setup_gp(M, L, noise_terms, gp_type == 1, w0);  // basis function
   // Rt
   real r_logmean = log(r_mean^2 / sqrt(r_sd^2 + r_mean^2));
   real r_logsd = sqrt(log(1 + (r_sd^2 / r_mean^2)));
@@ -41,38 +40,39 @@ transformed data{
   }
 }
 
-parameters{
+parameters {
   // gaussian process
-  array[fixed ? 0 : 1] real<lower = ls_min,upper=ls_max> rho;  // length scale of noise GP
-  array[fixed ? 0 : 1] real<lower = 0> alpha;    // scale of of noise GP
+  array[fixed ? 0 : 1] real<lower = ls_min, upper = ls_max> rho;  // length scale of noise GP
+  array[fixed ? 0 : 1] real<lower = 0> alpha;    // scale of noise GP
   vector[fixed ? 0 : M] eta;               // unconstrained noise
   // Rt
   vector[estimate_r] log_R;                // baseline reproduction number estimate (log)
-  array[estimate_r] real initial_infections ;    // seed infections
+  array[estimate_r] real initial_infections;    // seed infections
   array[estimate_r && seeding_time > 1 ? 1 : 0] real initial_growth; // seed growth rate
   array[bp_n > 0 ? 1 : 0] real<lower = 0> bp_sd; // standard deviation of breakpoint effect
   array[bp_n] real bp_effects;                   // Rt breakpoint effects
   // observation model
-
   vector<lower = delay_params_lower>[delay_params_length] delay_params; // delay parameters
-  simplex[week_effect] day_of_week_simplex;// day of week reporting effect
-  array[obs_scale_sd > 0 ? 1 : 0] real<lower = 0, upper = 1> frac_obs;     // fraction of cases that are ultimately observed
-  array[model_type] real<lower = 0> rep_phi;     // overdispersion of the reporting process
+  simplex[week_effect] day_of_week_simplex; // day of week reporting effect
+  array[obs_scale_sd > 0 ? 1 : 0] real<lower = 0, upper = 1> frac_obs; // fraction of cases that are ultimately observed
+  array[model_type] real<lower = 0> rep_phi; // overdispersion of the reporting process
 }
 
 transformed parameters {
-  vector[fixed ? 0 : noise_terms] noise;                    // noise  generated by the gaussian process
+  vector[fixed ? 0 : noise_terms] noise;                    // noise generated by the gaussian process
   vector<lower = 0>[estimate_r > 0 ? ot_h : 0] R; // reproduction number
   vector[t] infections;                                     // latent infections
   vector[ot_h] reports;                                     // estimated reported cases
   vector[ot] obs_reports;                                   // observed estimated reported cases
   vector[estimate_r * (delay_type_max[gt_id] + 1)] gt_rev_pmf;
+
   // GP in noise - spectral densities
   profile("update gp") {
     if (!fixed) {
-      noise = update_gp(PHI, M, L, alpha[1], rho[1], eta, gp_type);
+      noise = update_gp(PHI, M, L, alpha[1], rho[1], eta, gp_type, nu);
     }
   }
+
   // Estimate latent infections
   if (estimate_r) {
     profile("gt") {
@@ -102,6 +102,7 @@ transformed parameters {
       );
     }
   }
+
   // convolve from latent infections to mean of observations
   if (delay_id) {
     vector[delay_type_max[delay_id] + 1] delay_rev_pmf;
@@ -119,12 +120,14 @@ transformed parameters {
   } else {
     reports = infections[(seeding_time + 1):t];
   }
+
   // weekly reporting effect
   if (week_effect > 1) {
     profile("day of the week") {
       reports = day_of_week_effect(reports, day_of_week, day_of_week_simplex);
     }
   }
+
   // scaling of reported cases by fraction observed
   if (obs_scale) {
     profile("scale") {
@@ -133,6 +136,7 @@ transformed parameters {
       );
     }
   }
+
   // truncate near time cases to observed reports
   if (trunc_id) {
     vector[delay_type_max[trunc_id] + 1] trunc_rev_cmf;
@@ -161,13 +165,15 @@ model {
       );
     }
   }
-  // penalised priors for delay distributions
+
+  // penalized priors for delay distributions
   profile("delays lp") {
     delays_lp(
       delay_params, delay_params_mean, delay_params_sd, delay_params_groups,
       delay_dist, delay_weight
     );
   }
+
   if (estimate_r) {
     // priors on Rt
     profile("rt lp") {
@@ -177,12 +183,14 @@ model {
       );
     }
   }
+
   // prior observation scaling
   if (obs_scale_sd > 0) {
     profile("scale lp") {
       frac_obs[1] ~ normal(obs_scale_mean, obs_scale_sd) T[0, 1];
     }
   }
+
   // observed reports from mean of reports (update likelihood)
   if (likelihood) {
     profile("report lp") {
@@ -196,11 +204,12 @@ model {
 
 generated quantities {
   array[ot_h] int imputed_reports;
-  vector[estimate_r > 0 ? 0: ot_h] gen_R;
+  vector[estimate_r > 0 ? 0 : ot_h] gen_R;
   vector[ot_h - 1] r;
   vector[return_likelihood ? ot : 0] log_lik;
+
   profile("generated quantities") {
-    if (estimate_r == 0){
+    if (estimate_r == 0) {
       // sample generation time
       vector[delay_params_length] delay_params_sample = to_vector(normal_lb_rng(
         delay_params_mean, delay_params_sd, delay_params_lower
@@ -216,10 +225,13 @@ generated quantities {
         infections, seeding_time, sampled_gt_rev_pmf, rt_half_window
       );
     }
+
     // estimate growth from infections
     r = calculate_growth(infections, seeding_time + 1);
+
     // simulate reported cases
     imputed_reports = report_rng(reports, rep_phi, model_type);
+
     // log likelihood of model
     if (return_likelihood) {
       log_lik = report_log_lik(