improvements for gp() predictions and vis

R/dynamic.R

@@ -135,6 +135,9 @@ interpret_mvgam = function(formula, N){
 
   facs <- colnames(attr(terms.formula(formula), 'factors'))
 
+  # Check if formula has an intercept
+  keep_intercept <- attr(terms(formula), 'intercept') == 1
+
   # Re-arrange so that random effects always come last
   if(any(grepl('bs = \"re\"', facs, fixed = TRUE))){
     newfacs <- facs[!grepl('bs = \"re\"', facs, fixed = TRUE)]
@@ -263,5 +266,10 @@ interpret_mvgam = function(formula, N){
     updated_formula <- newformula
   }
 
+  if(!keep_intercept){
+    updated_formula <- update(updated_formula, . ~ . - 1)
+    attr(updated_formula, '.Environment') <- attr(newformula, '.Environment')
+  }
+
   return(updated_formula)
 }

R/gp.R

@@ -93,6 +93,7 @@ make_gp_additions = function(gp_details, data,
       coefs_replace[[x]] <- which_replace
     }
 
+  # Replace basis functions with gp() eigenfunctions
   newX <- model_data$X
 
   # Training data eigenfunctions
@@ -155,6 +156,22 @@ make_gp_additions = function(gp_details, data,
   # Add the GP attribute table to the mgcv_model
   attr(mgcv_model, 'gp_att_table') <- gp_att_table
 
+  # Assign GP labels to smooths
+  gp_assign <- data.frame(label = unlist(purrr::map(gp_att_table, 'name')),
+                          first.para = unlist(purrr::map(gp_att_table, 'first_coef')),
+                          last.para = unlist(purrr::map(gp_att_table, 'last_coef')),
+                          by = unlist(purrr::map(gp_att_table, 'by')))
+  for(i in seq_along(mgcv_model$smooth)){
+    if(mgcv_model$smooth[[i]]$label %in%
+       gsub('gp(', 's(', gp_assign$label, fixed = TRUE) &
+       mgcv_model$smooth[[i]]$first.para %in% gp_assign$first.para){
+      mgcv_model$smooth[[i]]$gp_term <- TRUE
+    } else {
+      mgcv_model$smooth[[i]]$gp_term <- FALSE
+    }
+  }
+
+  # Return
   return(list(model_data = model_data,
               mgcv_model = mgcv_model,
               gp_stan_lines = gp_stan_lines,
@@ -303,34 +320,37 @@ sim_hilbert_gp = function(alpha_gp,
 #' Mean-center and scale the particular covariate of interest
 #' so that the maximum Euclidean distance between any two points is 1
 #' @noRd
-scale_cov <- function(data, covariate, by, level, scale = TRUE,
+scale_cov <- function(data, covariate, by, level,
                       mean, max_dist){
   Xgp <- data[[covariate]]
   if(!is.na(by) &
      !is.na(level)){
       Xgp <- data[[covariate]][data[[by]] == level]
    }
 
-  if(is.na(mean)){
-    Xgp_mean <- mean(Xgp, na.rm = TRUE)
-  } else {
-    Xgp_mean <- mean
-  }
-
+  # Compute max Euclidean distance if not supplied
   if(is.na(max_dist)){
     Xgp_max_dist <- sqrt(max(brms:::diff_quad(Xgp)))
   } else {
     Xgp_max_dist <- max_dist
   }
 
-  if(scale){
-    # Mean center and divide by max euclidean distance
-    (Xgp - Xgp_mean) / Xgp_max_dist
+  # Scale
+  Xgp <- Xgp / Xgp_max_dist
 
+  # Compute mean if not supplied (after scaling)
+  if(is.na(mean)){
+    Xgp_mean <- mean(Xgp, na.rm = TRUE)
   } else {
-    # Just mean center
-    Xgp - Xgp_mean
+    Xgp_mean <- mean
   }
+
+  # Center
+  Xgp <- Xgp - Xgp_mean
+
+  return(list(Xgp = Xgp,
+              Xgp_mean = Xgp_mean,
+              Xgp_max_dist = Xgp_max_dist))
 }
 
 #' prep GP eigenfunctions
@@ -355,8 +375,7 @@ prep_eigenfunctions = function(data,
                               by = by,
                               level = level,
                               mean = mean,
-                              max_dist = max_dist,
-                              scale = scale)
+                              max_dist = max_dist)$Xgp
 
   # Construct matrix of eigenfunctions
   eigenfunctions <- matrix(NA, nrow = length(covariate_cent),
@@ -431,25 +450,24 @@ prep_gp_covariate = function(data,
   if(def_alpha == ''){
     def_alpha<- 'student_t(3, 0, 2.5);'
   }
+
   # Prepare the covariate
+  if(scale){
+    max_dist <- NA
+  } else {
+    max_dist <- 1
+  }
+
   covariate_cent <- scale_cov(data = data,
                               covariate = covariate,
-                              scale = scale,
                               by = by,
                               mean = NA,
-                              max_dist = NA,
+                              max_dist = max_dist,
                               level = level)
 
-  Xgp <- data[[covariate]]
-  if(!is.na(by) &
-     !is.na(level)){
-    Xgp <- data[[covariate]][data[[by]] == level]
-  }
-
-  covariate_mean <- mean(Xgp, na.rm = TRUE)
-  covariate_max_dist <- ifelse(scale,
-                               sqrt(max(brms:::diff_quad(Xgp))),
-                               1)
+  covariate_mean <- covariate_cent$Xgp_mean
+  covariate_max_dist <- covariate_cent$Xgp_max_dist
+  covariate_cent <- covariate_cent$Xgp
 
   # Construct vector of eigenvalues for GP covariance matrix; the
   # same eigenvalues are always used in prediction, so we only need to
@@ -469,9 +487,10 @@ prep_gp_covariate = function(data,
                                         covariate = covariate,
                                         by = by,
                                         level = level,
+                                        L = L,
                                         k = k,
                                         boundary = boundary,
-                                        mean = NA,
+                                        mean = covariate_mean,
                                         max_dist = covariate_max_dist,
                                         scale = scale,
                                         initial_setup = TRUE)

R/mvgam.R

@@ -652,7 +652,7 @@ mvgam = function(formula,
     # in the model.frame
     formula <- gp_to_s(formula)
     if(!keep_intercept){
-      formula <- update(formula, trend_y  ~ . -1)
+      formula <- update(formula, . ~ . - 1)
     }
   }
 

R/plot.mvgam.R

@@ -134,12 +134,6 @@ plot.mvgam = function(x, type = 'residuals',
                  mgcv_plottable = object2$mgcv_model$smooth[[x]]$plot.me)
     }))
 
-    # Filter out any GP terms
-    if(!is.null(attr(object2$mgcv_model, 'gp_att_table'))){
-      gp_names <- unlist(purrr::map(attr(object2$mgcv_model, 'gp_att_table'), 'name'))
-      smooth_labs %>%
-        dplyr::filter(!label %in% gsub('gp(', 's(', gp_names, fixed = TRUE)) -> smooth_labs
-    }
     n_smooths <- NROW(smooth_labs)
     if(n_smooths == 0) stop("No smooth terms to plot. Use plot_predictions() to visualise other effects",
                             call. = FALSE)

R/plot_mvgam_smooth.R

@@ -115,17 +115,6 @@ plot_mvgam_smooth = function(object,
     smooth_int <- smooth
   }
 
-  # Check whether this is actually a gp() term
-  if(!is.null(attr(object2$mgcv_model, 'gp_att_table'))){
-    gp_names <- unlist(purrr::map(attr(object2$mgcv_model, 'gp_att_table'), 'name'))
-    if(any(grepl(object2$mgcv_model$smooth[[smooth_int]]$label,
-          gsub('gp(', 's(', gp_names, fixed = TRUE),
-          fixed = TRUE))){
-      stop(smooth, ' is a gp() term. Use plot_predictions() instead to visualise',
-           call. = FALSE)
-    }
-  }
-
   # Check whether this type of smooth is even plottable
   if(!object2$mgcv_model$smooth[[smooth_int]]$plot.me){
     stop(paste0('unable to plot ', object2$mgcv_model$smooth[[smooth_int]]$label,
@@ -290,14 +279,61 @@ plot_mvgam_smooth = function(object,
     # If this term has a by variable, need to use mgcv's plotting utilities
     if(object2$mgcv_model$smooth[[smooth_int]]$by != "NA"){
 
-      # Deal with by variables
-      by <- rep(1,length(pred_vals)); dat <- data.frame(x = pred_vals, by = by)
-      names(dat) <- c(object2$mgcv_model$smooth[[smooth_int]]$term,
-                      object2$mgcv_model$smooth[[smooth_int]]$by)
+      # Check if this is a gp() term
+      gp_term <- FALSE
+      if(!is.null(attr(object2$mgcv_model, 'gp_att_table'))){
+        gp_term <- object2$mgcv_model$smooth[[smooth_int]]$gp_term
+      }
+
+      if(gp_term){
+        object2$mgcv_model$smooth[[smooth_int]]$label <-
+          gsub('s(', 'gp(',
+               object2$mgcv_model$smooth[[smooth_int]]$label,
+               fixed = TRUE)
+        # Check if this is a factor by variable
+        is_fac <- is.factor(object2$obs_data[[object2$mgcv_model$smooth[[smooth_int]]$by]])
+
+        if(is_fac){
+          fac_levels <- levels(object2$obs_data[[object2$mgcv_model$smooth[[smooth_int]]$by]])
+          whichlevel <- vector()
+          for(i in seq_along(fac_levels)){
+            whichlevel[i] <- grepl(fac_levels[i], object2$mgcv_model$smooth[[smooth_int]]$label,
+                  fixed = TRUE)
+          }
+
+          pred_dat[[object2$mgcv_model$smooth[[smooth_int]]$by]] <-
+            rep(fac_levels[whichlevel], length(pred_dat$series))
+        }
+
+        if(!is_fac){
+          pred_dat[[object2$mgcv_model$smooth[[smooth_int]]$by]] <-
+            rep(1, length(pred_dat$series))
+        }
+
+        if(trend_effects){
+          Xp_term <- trend_Xp_matrix(newdata = pred_dat,
+                                trend_map = object2$trend_map,
+                                mgcv_model = object2$trend_mgcv_model)
+        } else {
+          Xp_term <- obs_Xp_matrix(newdata = pred_dat,
+                              mgcv_model = object2$mgcv_model)
+        }
+        Xp[,object2$mgcv_model$smooth[[smooth_int]]$first.para:
+             object2$mgcv_model$smooth[[smooth_int]]$last.para] <-
+          Xp_term[,object2$mgcv_model$smooth[[smooth_int]]$first.para:
+                    object2$mgcv_model$smooth[[smooth_int]]$last.para]
+
+      } else {
+        # Deal with by variables in non-gp() smooths
+        by <- rep(1,length(pred_vals)); dat <- data.frame(x = pred_vals, by = by)
+        names(dat) <- c(object2$mgcv_model$smooth[[smooth_int]]$term,
+                        object2$mgcv_model$smooth[[smooth_int]]$by)
+
+        Xp_term <- mgcv::PredictMat(object2$mgcv_model$smooth[[smooth_int]], dat)
+        Xp[,object2$mgcv_model$smooth[[smooth_int]]$first.para:
+             object2$mgcv_model$smooth[[smooth_int]]$last.para] <- Xp_term
+      }
 
-      Xp_term <- mgcv::PredictMat(object2$mgcv_model$smooth[[smooth_int]], dat)
-      Xp[,object2$mgcv_model$smooth[[smooth_int]]$first.para:
-           object2$mgcv_model$smooth[[smooth_int]]$last.para] <- Xp_term
     }
 
     # Extract GAM coefficients