R/simData.R

In devoges/fstat: fstat

#' @useDynLib fstat, .registration = TRUE
#' @importFrom Rcpp sourceCpp
NULL

sim_data <- function(n,
                     p,
                     beta,
                     Z_ks = NULL,
                     Z_sigmas = NULL,
                     S_n = NULL,
                     S_sigmas = NULL,
                     sigE = 1) {
  X <- matrix(rnorm(n * p), n, p)
  X <- cbind(1, X) %>% `colnames<-`(c("(Intercept)", "X_" %% 1:p))
  beta <- matrix(beta, ncol = 1)
  Xbeta <- X %*% beta

  Y <- Xbeta + rnorm(n, sd = sqrt(sigE))

  out <- list()
  out$Y <- Y
  out$X <- X

  if ( !is.null(Z_ks) ) {
    message("doing Zs")

    Zs_names <- "Z_" %% 1:length(Z_ks)
    names(Z_ks) <- Zs_names
    names(Z_sigmas) <- Zs_names

    Zs <- nlapply(Zs_names, function(x) {
      matrix(rnorm(n * Z_ks[x]), n, Z_ks[x])
    })
    Us <-  nlapply(Zs_names, function(x) {
      mvrnorm(n = 1,
              mu = rep(0, Z_ks[x]),
              Sigma = diag(Z_ks[x]) * Z_sigmas[x])
    })

    ZUs <- nlapply(Zs_names, function(x) {
      Zs[[x]] %*% Us[[x]]
    })

    ZU <- Reduce(`+`, ZUs)

    Y <- Y + ZU

    out$Zs <- Zs

  }

  if ( !is.null(S_n) & S_n > 0) {

    Ss_names <- "S_" %% 1:S_n
    names(S_sigmas) <- Ss_names

    Ss <- nlapply(Ss_names, function(x) {
      tcrossprod(matrix(rnorm(n^2), n, n))
    })

    SUs <- nlapply(Ss_names, function(x) {
      mvrnorm(n = 1,
              mu = rep(0, n),
              Sigma = diag(n) * S_sigmas[x])
    })

    SU <- Reduce(`+`, SUs)

    Y <- Y + SU

    out$Ss <- Ss
  }

  return(out)
}

# debugonce(sim_data)
# d <- sim_data(n = 1000,
#               p = 2,
#               beta = c(1, 0, 1),
#               Z_ks = c(200, 1),
#               Z_sigmas = c(0.5, 0.5),
#               S_n = 0,
#               S_sigmas = NULL)
# str(d)