R/generate.R

In koekvall/mpredcc: Fully Likelihood-based Principal Components Regression

#' Generate data from Mean and Predictor Covariance-Connected (MPREDCC)
#' regression model; iid or autoregressive structures supported.
#'
#' @param n The number of observations to generate.
#' @param alpha The p-vector of coefficients such that beta = Psi x alpha.
#' @param Psi A p x p SPSD matrix of rank k giving the predictor covariance
#'            structure.
#' @param ssy Noise variance for response
#' @param ssx Noise variance for predictors
#' @param mu_Y Mean parameter for responses
#' @param mu_X Mean parameter for predictors
#' @param a Vector of autoregressive parameters for the response, first element
#'          is first lag
#' @param A Matrix of autoregressive parameters for the predictors, first p rows
#'          is for the first lag
#' @param Y_start Starting values if autoregressive responses; default to zeros
#' @param X_start Starting values if autoregressive predictors; default to zeros
#' @return List with responses (Y) and predictors (X)
#' @export
#'                
generate_mpredcc <- function(n = 100, alpha, Psi, ssy = 1, ssx = 1, mu_Y = 0,
                             mu_X = rep(0, ncol(Psi)),
                             a = NULL, A = NULL, Y_start = NULL,
                             X_start = NULL)
{
  p <- ncol(Psi)
  
  if(!is.null(a)){
    q_Y <- length(a)
  } else{
    q_Y <- 0
  }
  
  if(!is.null(A)){
    q_X <- nrow(A) / p
  } else{
    q_X <- 0
  }
  
  Y <- rep(0, n + q_Y)
  if(!is.null(Y_start) & q_Y > 0){
    Y[1:q_Y] <- Y_start
  }
  
  X <- matrix(0, nrow = n + q_X, ncol = p)
  if(!is.null(X_start) & q_X > 0){
    X[1:q_X, ] <- X_start
  }
  # Cholesky decomposition of predictor covariance
  R <- chol(Psi + diag(ssx, p))
  
  for(ii in 1:n){
    idx_X <- ii + q_X
    X[idx_X, ] <- stats::rnorm(p) %*% R + mu_X
    if(!is.null(A) & q_X > 0){
      X[idx_X, ] <- X[idx_X, ] + crossprod(A, as.vector(t(X[(idx_X - 1):(idx_X - q_X), ])))
    }
    
    idx_Y <- ii + q_Y
    Y[idx_Y] <- stats::rnorm(n = 1, mean = mu_Y, sd = sqrt(ssy)) + X[idx_X, ] %*% (Psi %*% alpha)
    if(!is.null(a) & q_Y > 0){
      Y[idx_Y] <- Y[idx_Y] + sum(Y[(idx_Y - 1):(idx_Y - q_Y)] * a)
    }
  }
  return(list(Y = Y, X = X))
}