R/build_dllm.R

In dlmwwbe: Dynamic Linear Model for Wastewater-Based Epidemiology

#' Build a Dynamic Local Level Model (DLLM)
#'
#' Constructs a dynamic linear model (DLLM) object using the dlm package.
#'
#' @param params Numeric vector of parameters (variances only, must be positive). If \code{V} or \code{W} are not provided,
#'   \code{params} is used to construct them.
#' @param obs_cols Character vector specifying the names of observation columns.
#' @param S Integer. The number of latent states.
#' @param prior A list of prior specifications. Default priors are provided if not overridden.
#' @param cov_structure Character; one of \code{"diagonal-equal"} or \code{"diagonal-unequal"}.
#' @param Ft Optional observation matrix. If \code{NULL}, a default matrix is constructed.
#' @param Gt Optional transition matrix. If \code{NULL}, an identity matrix is used.
#' @param V Optional observation covariance matrix.
#' @param W Optional system covariance matrix.
#' @param X0 a vector of the first row of the data
#'
#' @return A DLM object (from \code{dlm::dlm}) with an additional attribute \code{"model_info"}.
#'
#'
#' @noRd
build_dllm <- function(params = NULL,
                      x0,
                      nW,
                      nV,
                      obs_cols,
                      S,
                      k,
                      prior = list(),
                      equal.state.var = FALSE,
                      equal.obs.var = FALSE
                      ) {
  # Ft - observation matrix
  if (S == 'univariate'){
    Ft <- as.matrix(rep(x=1, times=k), ncol=1)
  }else{
    Ft <- diag(x=1, nrow=k)
  }
  # Gt - transition matrix
  if (S == 'univariate'){
    Gt <- diag(x=1, nrow=1)
  }else{
    Gt <- diag(x=1, nrow=k)
  }

  # Wt & Vt - observation and state covariance matrice
  dimW <- ifelse(S=='univariate', 1, k)


  Wt <- diag(x=exp(params[c(1:nW)]), nrow=dimW )
  Vt <- diag(x=exp(params[-c(1:nW)]), nrow=k)

  # default prior
  if(S=='univariate'){
    m0 <-  mean(x0, na.rm=TRUE)
  }else{
    m0 <- x0
  }
  #m0 <- rep(x=0, times=dimW)
  C0 <- diag(x=10**-6, nrow=dimW)
  prior <- utils::modifyList(list(m0=m0, C0=C0), prior)

  # Construct the DLM using the dlm package.
  mod <- dlm::dlm(
    m0 = prior$m0,
    C0 = prior$C0,
    FF = Ft,
    GG = Gt,
    V = Vt,
    W = Wt
  )

  # Attach additional model info as an attribute.
  attr(mod, "model_info") <- list(
    obs_cols = obs_cols,
    S = S,
    equal.state.var=equal.state.var,
    equal.obs.var=equal.obs.var
  )

  return(mod)
}