R/design-matrix.R

In bivarhr: Bivariate Hurdle Regression with Bayesian Model Averaging

#' Create Lag Matrix
#'
#' Creates a matrix of lagged values for a numeric vector.
#'
#' @param x Numeric vector.
#' @param k Integer; maximum lag order.
#'
#' @return A matrix with k columns containing lags 1 through k.
#'
#' @export

make_lags <- function(x, k) {
  if (k <= 0) return(matrix(numeric(0), nrow=length(x), ncol=0))
  out <- sapply(1:k, function(j) dplyr::lag(x, j))
  colnames(out) <- paste0(deparse(substitute(x)), "_L", 1:k)
  out
}

#' Build Design Matrices for Bivariate Hurdle Model
#'
#' Constructs design matrices for the zero and count components
#' of both outcome variables with cross-lags, trends, regimes,
#' transition dummies, and control variables.
#'
#' @param DT A data.table with required columns.
#' @param k Integer; lag order.
#' @param include_C_to_I Logical; include C lags in I equations.
#' @param include_I_to_C Logical; include I lags in C equations.
#' @param include_trend Logical; include polynomial time trend.
#' @param controls Character vector of control variable names.
#' @param include_regimes Logical; include regime dummies.
#' @param include_transitions Logical; include transition dummies.
#'
#' @return A list containing design matrices and outcome vectors.
#'
#' @export

build_design <- function(DT, k,
                         include_C_to_I = TRUE,
                         include_I_to_C = TRUE,
                         include_trend = TRUE,
                         controls = character(0),
                         include_regimes = TRUE,
                         include_transitions = TRUE) {
  stopifnot(k >= 0)
  X_pi_I <- X_mu_I <- X_pi_C <- X_mu_C <- NULL
  nlags_pi_I <- nlags_mu_I <- nlags_pi_C <- nlags_mu_C <- 0L
  
  if (include_C_to_I && k>0) {
    X_pi_I <- make_lags(DT$zC, k); colnames(X_pi_I) <- paste0("zC_L",1:k)
    X_mu_I <- make_lags(DT$C,  k); colnames(X_mu_I) <- paste0("C_L",1:k)
    nlags_pi_I <- ncol(X_pi_I); nlags_mu_I <- ncol(X_mu_I)
  }
  if (include_I_to_C && k>0) {
    X_pi_C <- make_lags(DT$zI, k); colnames(X_pi_C) <- paste0("zI_L",1:k)
    X_mu_C <- make_lags(DT$I,  k); colnames(X_mu_C) <- paste0("I_L",1:k)
    nlags_pi_C <- ncol(X_pi_C); nlags_mu_C <- ncol(X_mu_C)
  }
  
  trend_mat <- if (include_trend) cbind(t_norm=DT$t_norm, t_poly2=DT$t_poly2) else NULL
  regimes_mat <- if (include_regimes) model.matrix(~ Regime, DT)[,-1, drop=FALSE] else NULL
  trans_mat <- if (include_transitions) as.matrix(DT[, .(trans_PS, trans_SF, trans_FC)]) else NULL
  if (!is.null(trans_mat)) colnames(trans_mat) <- c("trans_PS","trans_SF","trans_FC")
  
  ctrl_mat <- NULL
  if (length(controls) > 0) {
    ctrl_mat <- if (length(controls) > 0) as.matrix(DT[, ..controls]) else NULL
    colnames(ctrl_mat) <- controls
  }
  
  bindX <- function(Xlags, trend, regimes, trans, ctrl) {
    X <- cbind(Xlags, trend, regimes, trans, ctrl)
    if (is.null(X)) matrix(0, nrow=nrow(DT), ncol=0) else X
  }
  
  X_pi_I <- bindX(X_pi_I, trend_mat, regimes_mat, trans_mat, ctrl_mat)
  X_mu_I <- bindX(X_mu_I, trend_mat, regimes_mat, trans_mat, ctrl_mat)
  X_pi_C <- bindX(X_pi_C, trend_mat, regimes_mat, trans_mat, ctrl_mat)
  X_mu_C <- bindX(X_mu_C, trend_mat, regimes_mat, trans_mat, ctrl_mat)
  
  idx <- (k+1):nrow(DT)
  list(
    idx = idx,
    y_I = DT$I[idx], y_C = DT$C[idx],
    z_I = DT$zI[idx], z_C = DT$zC[idx],
    log_exposure50 = DT$log_exposure50[idx],
    X_pi_I = X_pi_I[idx,,drop=FALSE],
    X_mu_I = X_mu_I[idx,,drop=FALSE],
    X_pi_C = X_pi_C[idx,,drop=FALSE],
    X_mu_C = X_mu_C[idx,,drop=FALSE],
    nlags_pi_I = nlags_pi_I, nlags_mu_I = nlags_mu_I,
    nlags_pi_C = nlags_pi_C, nlags_mu_C = nlags_mu_C
  )
}