Smooth construct, predict, and plot for random effects with numeric inputs

smooth_re2.R

#' Create a random effect basis with integers rather than factors
#' @import mgcv
#' @export
smooth.construct.re2.smooth.spec <- function (object, data, knots) {
  if (!is.null(object$id))
    stop("random effects don't work with ids.")
  if(any(sapply(data, is.numeric))) data <- lapply(data, as.factor)   ## <-- All I did was this (and below)
  form <- as.formula(paste("~", paste(object$term, collapse = ":"),
                           "-1"))
  object$X <- model.matrix(form, data)
  object$bs.dim <- ncol(object$X)
  if (inherits(object, "tensor.smooth.spec")) {
    object$margin <- list()
    maxd <- maxi <- 0
    for (i in 1:object$dim) {
      form1 <- as.formula(paste("~", object$term[i], "-1"))
      object$margin[[i]] <- list(X = model.matrix(form1,
                                                  data), term = object$term[i], form = form1,
                                 by = "NA")
      class(object$margin[[i]]) <- "random.effect2"
      d <- ncol(object$margin[[i]]$X)
      if (d > maxd) {
        maxi <- i
        maxd <- d
      }
    }
    if (maxi < object$dim) {
      ns <- object$dim
      ind <- 1:ns
      ind[maxi] <- ns
      ind[ns] <- maxi
      object$margin <- object$margin[ind]
      object$term <- rep("", 0)
      for (i in 1:ns) object$term <- c(object$term, object$margin[[i]]$term)
      object$label <- paste0(substr(object$label, 1, 2),
                             paste0(object$term, collapse = ","), ")", collapse = "")
      object$rind <- ind
      if (!is.null(object$xt$S))
        stop("Please put term with most levels last in 're' to avoid spoiling supplied penalties")
    }
  }
  if (is.null(object$xt$S)) {
    object$S <- list(diag(object$bs.dim))
    object$rank <- object$bs.dim
  }
  else {
    object$S <- if (is.list(object$xt$S))
      object$xt$S
    else list(object$xt$S)
    for (i in 1:length(object$S)) {
      if (ncol(object$S[[i]]) != object$bs.dim || nrow(object$S[[i]]) !=
          object$bs.dim)
        stop("supplied S matrices are wrong diminsion")
    }
    object$rank <- object$xt$rank
  }
  object$null.space.dim <- 0
  object$C <- matrix(0, 0, ncol(object$X))
  object$form <- form
  object$side.constrain <- FALSE
  object$plot.me <- TRUE
  object$te.ok <- if (inherits(object, "tensor.smooth.spec"))
    0
  else 2
  object$random <- TRUE
  object$noterp <- TRUE
  class(object) <- if (inherits(object, "tensor.smooth.spec"))
    c("random.effect2", "tensor.smooth")
  else "random.effect2"
  object
}

#' Predict from a random effect basis with integers rather than factors
#' @import mgcv
#' @export
Predict.matrix.random.effect2 <- function (object, data) {
  if(any(sapply(data, is.numeric))) data <- lapply(data, as.factor)  # <-- And did it again here
  X <- model.matrix(object$form, model.frame(object$form, data,
                                             na.action = na.pass))
  X[!is.finite(X)] <- 0
  X
}

#' Plot a numeric random effect
#' @import mgcv
#' @export
plot.random.effect2 <- function (x, P = NULL, data = NULL, label = "", se1.mult = 1,
                                 se2.mult = 2, partial.resids = FALSE, rug = TRUE, se = TRUE,
                                 scale = -1, n = 100, n2 = 40, n3 = 3, pers = FALSE, theta = 30,
                                 phi = 30, jit = FALSE, xlab = NULL, ylab = NULL, main = NULL,
                                 ylim = NULL, xlim = NULL, too.far = 0.1, shade = FALSE,
                                 shade.col = "gray80", shift = 0, trans = I, by.resids = FALSE,
                                 scheme = 0, ...) {
  # Did nothing here at all, just renamed the function to `....random.effect2`
  if (is.null(P)) {
    if (!x$plot.me)
      return(NULL)
    else {
      raw <- data[x$term][[1]]
      p <- x$last.para - x$first.para + 1
      X <- diag(p)
      if (is.null(xlab))
        xlabel <- "Gaussian quantiles"
      else xlabel <- xlab
      if (is.null(ylab))
        ylabel <- "effects"
      else ylabel <- ylab
      if (!is.null(main))
        label <- main
      return(list(X = X, scale = FALSE, se = FALSE, raw = raw,
                  xlab = xlabel, ylab = ylabel, main = label))
    }
  }
  else {
    b <- as.numeric(trans(P$fit + shift))
    qqnorm(b, main = P$main, xlab = P$xlab, ylab = P$ylab,
           ...)
    qqline(b)
  }
}