R/mc_mec.R

In mcmcsae: Markov Chain Monte Carlo Small Area Estimation

#' Create a model component object for a regression (fixed effects) component
#' in the linear predictor with measurement errors in quantitative covariates
#'
#' This function is intended to be used on the right hand side of the
#' \code{formula} argument to \code{\link{create_sampler}} or
#' \code{\link{generate_data}}. It creates an additive regression term in the
#' model's linear predictor. Covariates are assumed to be measured subject
#' to normally distributed errors with zero mean and variance specified using
#' the \code{formula} or \code{V} arguments. Note that this means that \code{formula}
#' should only contain quantitative variables, and no intercept.
#' By default, the prior for the regression
#' coefficients is improper uniform. A proper normal prior can be set up
#' using function \code{\link{pr_normal}}, and passed to argument \code{prior}.
#' It should be noted that \code{\link{pr_normal}} expects a precision matrix
#' as input for its second argument, and that the prior variance (matrix) is
#' taken to be the inverse of this precision matrix, where in case the
#' model's family is \code{"gaussian"} this matrix is additionally
#' multiplied by the residual scalar variance parameter \code{sigma_^2}.
#'
#' @examples
#' \donttest{
#' # example of Ybarra and Lohr (2008)
#' m <- 50
#' X <- rnorm(m, mean=5, sd=3)  # true covariate values
#' v <- rnorm(m, sd=2)
#' theta <- 1 + 3*X + v  # true values
#' psi <- rgamma(m, shape=4.5, scale=2)
#' e <- rnorm(m, sd=sqrt(psi))  # sampling error
#' y <- theta + e  # direct estimates
#' C <- c(rep(3, 10), rep(0, 40))  # measurement error for first 10 values
#' W <- X + rnorm(m, sd=sqrt(C))  # covariate subject to measurement error
#'
#' # fit Ybarra-Lohr model
#' sampler <- create_sampler(
#'   y ~ 1 + mec(~ 0 + W, V=C, name="ME") + gen(factor=~local_),
#'   family=f_gaussian(var.prior=pr_fixed(1), var.vec=~psi),
#'   linpred="fitted"
#' )
#' sim <- MCMCsim(sampler, n.iter=800, n.chain=2, store.all=TRUE, verbose=FALSE)
#' (summ <- summary(sim))
#' plot(X, W, xlab="true X", ylab="inferred X")
#' points(X, summ$ME_X[, "Mean"], col="green")
#' abline(0, 1, col="red")
#' legend("topleft", legend=c("prior mean", "posterior mean"), col=c("black", "green"), pch=c(1,1))
#' }
#'
#' @export
#' @param formula a formula specifying the predictors subject to measurement error
#'  and possibly their variances as well. In the latter case the formula syntax
#'  \code{~ (x1 | V.x1) + (x2 | V.x2) + ...} should be used where \code{x1, x2, ...}
#'  are the names of (quantitative) predictors and \code{V.x1, V.x2, ...} are the names
#'  of the variables holding the corresponding measurement error variances.
#'  If only the predictors are specified
#'  the formula has the usual form \code{~ x1 + x2 + ...}. In that case variances
#'  should be specified using argument \code{V}.
#'  All variable names are looked up in the data frame
#'  passed as \code{data} argument to \code{\link{create_sampler}} or
#'  \code{\link{generate_data}}, or in \code{environment(formula)}.
#' @param sparse whether the model matrix associated with \code{formula} should
#'  be sparse. The default is to base this on a simple heuristic.
#' @param X a (possibly sparse) design matrix can be specified directly, as an
#'  alternative to the creation of one based on \code{formula}. If \code{X} is
#'  specified \code{formula} is ignored.
#' @param V measurement error variance; can contain zeros
# TODO in general case this can be a list of q x q precision matrices where q
#      is the number of covariates
#' @param prior prior specification for the regression coefficients. Currently only
#'  normal priors are supported, specified using function \code{\link{pr_normal}}.
#' @param scale.prior whether the normal prior's scale includes the model's
#'  global residual scale as a factor. By default \code{TRUE}, but only relevant
#'  for a proper normal prior specified by \code{pr_normal} for families
#'  \code{"gaussian"}, \code{"student-t"} and \code{"gaussian-gamma"}.
#' @param constraints optional linear equality and/or inequality constraints
#'  imposed on the vector of regression coefficients. Use function
#'  \code{\link{set_constraints}} to specify the constraint matrices and
#'  right-hand sides.
#' @param Q0 prior precision matrix for the regression effects. The default is a
#'  zero matrix corresponding to a noninformative improper prior.
#'  It can be specified as a scalar value, as a numeric vector of appropriate
#'  length, or as a matrix object. DEPRECATED, please use argument \code{prior}
#'  instead, i.e. \code{prior = pr_normal(mean = b0.value, precision = Q0.value)}.
#' @param b0 prior mean for the regression effect. Defaults to a zero vector.
#'  It can be specified as a scalar value or as a numeric vector of
#'  appropriate length. DEPRECATED, please use argument \code{prior}
#'  instead, i.e. \code{prior = pr_normal(mean = b0.value, precision = Q0.value)}.
#' @param name the name of the model component. This name is used in the output of the
#'  MCMC simulation function \code{\link{MCMCsim}}. By default the name will be 'reg'
#'  with the number of the model term attached.
#' @param debug if \code{TRUE} a breakpoint is set at the beginning of the posterior
#'  draw function associated with this model component. Mainly intended for developers.
#' @returns An object with precomputed quantities and functions for sampling from
#'  prior or conditional posterior distributions for this model component,
#'  intended for internal use by other package functions.
#' @references
#'  L.M. Ybarra and S.L. Lohr (2008).
#'    Small area estimation when auxiliary information is measured with error.
#'    Biometrika 95(4), 919-931.
#'
#'  S. Arima, G.S. Datta and B. Liseo (2015).
#'    Bayesian estimators for small area models when auxiliary information is measured with error.
#'    Scandinavian Journal of Statistics 42(2), 518-529.
mec <- function(formula = ~ 1, sparse=NULL, X=NULL, V=NULL,
                prior=NULL, scale.prior=TRUE, constraints=NULL,
                Q0=NULL, b0=NULL,  # deprecated
                name="", debug=FALSE) {
  stop("function 'mec' should only be used inside a formula")
}

mc_mec <- function(formula = ~ 1, sparse=NULL, X=NULL, V=NULL,
                   prior=NULL, scale.prior=TRUE, constraints=NULL,
                   Q0=NULL, b0=NULL,
                   name="", debug=FALSE,
                   sc, fam, in.block, prior.only, compute.weights=FALSE, data) {
  type <- "mec"
  if (name == "") stop("missing model component name")
  store.default <- name

  if (fam[["family"]] == "multinomial") stop("TBI: 'mec' term in multinomial model")
  if (fam[["family"]] == "multi") stop("TBI: 'mec' term in multi response family model")
  if (!scale.prior) stop("scale.prior=FALSE not yet supported for mec component")
  
  if (is.null(X)) {
    # TODO: warn if formula contains explicit intercept or categorical terms
    vs <- as.list(attr(terms(formula), "variables")[-1L])
    if (allv(lengths(vs), 3L)) {
      if (!all(b_apply(vs, \(x) identical(x[[1L]], as.name("|"))))) stop("invalid 'formula' argument")
      V.in.formula <- TRUE
      formula.X <- as.formula(paste0("~ 0 + ", paste0(s_apply(vs, \(x) deparse(x[[2L]])), collapse=" + ")), env=environment(formula))
      if (grepl("|", as.character(formula.X)[2L], fixed=TRUE)) stop("invalid 'formula'")  # maybe '()' forgotten?
      formula.V <- as.formula(paste0("~ 0 + ", paste0(s_apply(vs, \(x) deparse(x[[3L]])), collapse=" + ")), env=environment(formula))
      X <- model_matrix(formula.X, data, sparse=sparse)
      V <- model_matrix(formula.V, data, sparse=sparse)
    } else {
      if (all(lengths(vs) <= 2L)) {
        V.in.formula <- FALSE
        formula.X <- as.formula(paste0("~ 0 + ", paste0(s_apply(vs, deparse), collapse=" + ")), env=environment(formula))
        X <- model_matrix(formula.X, data, sparse=sparse)
      } else {
        stop("invalid 'formula' argument")
      }
    }
  } else {
    if (is.null(dimnames(X)[[2L]])) colnames(X) <- seq_len(ncol(X))
    V.in.formula <- FALSE
  }
  if (nrow(X) != fam[["n"]]) stop("design matrix with incompatible number of rows")
  coef.names <- dimnames(X)[[2L]]
  label_funs <- list()
  label_funs[[name]] <- function() coef.names
  X <- economizeMatrix(X, sparse=sparse, strip.names=TRUE, check=TRUE)
  q <- ncol(X)

  if (!V.in.formula) {
    if (is.null(V)) stop("no measurement error variances supplied in mec component")
    if (is.vector(V)) {
      if (q != 1L) stop("'V' should be a matrix")
      V <- matrix(V, ncol=1L)
    }
    if (nrow(V) != fam[["n"]] || ncol(V) != q) stop("wrong size of variance matrix 'V'")
    # TODO allow n x q matrix (uncorrelated measurement errors), or list of n q x q symmetric matrices
  }
  # determine units with measurement error
  V <- unname(V)
  if (any(V < 0)) stop("negative measurement error variance(s)")
  i.me <- whichv(dapply(V, \(x) all(x > 0), MARGIN=1L), TRUE)  # TODO add tolerance + allow list of matrix
  nme <- length(i.me)
  if (nme == fam[["n"]]) i.me <- seq_len(fam[["n"]])
  if (q == 1L) {
    QME <- 1 / V[i.me, ]
    rm(V)
  } else {
    V <- V[i.me, ]
  }

  if (fam[["Q0.type"]] == "unit") {
    if (q == 1L) v0 <- 1 else v0 <- rep.int(1, nme)
  } else {
    if (fam[["Q0.type"]] == "symm") stop("not supported: measurement error component and non-diagonal sampling variance")
    v0 <- 1/diag(fam[["Q0"]])[i.me]
  }

  if (!is.null(b0) || !is.null(Q0)) {
    warn("arguments 'b0' and 'Q0' are deprecated; please use argument 'prior' instead")
    if (is.null(prior)) prior <- pr_normal(mean = if (is.null(b0)) 0 else b0, precision = if (is.null(Q0)) 0 else Q0)
  } else {
    if (is.null(prior)) prior <- pr_normal(mean=0, precision=0)
  }
  if (prior[["type"]] != "normal") stop("only a normal prior is currently supported for 'mec' model component effects")
  use.sigma <- !fam[["sigma.fixed"]] && scale.prior
  prior$init(q, coef.names, sparse=if (in.block) TRUE else NULL, sigma=!fam[["sigma.fixed"]])  # mc_block uses x-slot of precision matrix
  informative.prior <- prior[["informative"]]
  Q0 <- prior[["precision"]]
  zero.mean <- !informative.prior || allv(prior[["mean"]], 0)
  if (zero.mean) {
    prior$mean <- 0
    Q0b0 <- 0
  } else {
    if (length(prior[["mean"]]) == 1L)
      Q0b0 <- Q0 %m*v% rep.int(prior[["mean"]], q)
    else
      Q0b0 <- Q0 %m*v% prior[["mean"]]
  }
  is.proper <- informative.prior

  if (!zero.mean && compute.weights) stop("weights cannot be computed if some coefficients have non-zero prior means")

  name_X <- paste0(name, "_X")
  lp <- function(p) p[[name_X]] %m*v% p[[name]]
  lp_update <- function(x, plus=TRUE, p) mv_update(x, plus, p[[name_X]], p[[name]])
  # draws_linpred method acts on (subset of) mcdraws object, used in fitted() and pointwise log-likelihood llh_i functions
  draws_linpred <- function(obj, units=NULL, chains=NULL, draws=NULL, matrix=FALSE) {
    Xi <- get_from(obj[[name_X]], vars=units, chains=chains, draws=draws)
    if (matrix) {
      out <- tcrossprod(as.matrix.dc(get_from(obj[[name]], chains=chains, draws=draws), colnames=FALSE), Xi)
    } else {
      out <- list()
      for (ch in seq_along(chains))
        out[[ch]] <- tcrossprod(get_from(obj[[name]], chains=chains[ch], draws=draws)[[1L]], Xi)
    }
    out
  }

  # function that creates an oos prediction function (closure) based on new data
  # this computes the contribution of this component to the linear predictor
  make_predict <- function(newdata=NULL, Xnew=NULL, verbose=TRUE) {
    if (is.null(newdata)) {
      if (is.null(Xnew)) {
        # in-sample prediction
        Xnew <- X
        linpred <- function(p) p[[name_X]] %m*v% p[[name]]
        linpred_update <- function(x, plus=TRUE, p) mv_update(x, plus, p[[name_X]], p[[name]])
      } else {
        stop("for out-of-sample prediction with 'mec' components a data frame ",
             "must be supplied")
      }
    } else {
      if (!V.in.formula) stop("for out-of-sample prediction with 'mec' components make sure to specify ",
                              "the measurement error variances using the formula argument")
      Xnew <- model_matrix(formula.X, newdata, sparse=sparse)
      dimnames(Xnew) <- list(NULL, NULL)
      SEnew <- model_matrix(formula.V, newdata, sparse=sparse)
      dimnames(SEnew) <- list(NULL, NULL)
      if (any(SEnew < 0)) stop("negative measurement error variance(s) in 'newdata'")
      SEnew <- sqrt(SEnew)
      if (nrow(SEnew) < 0.5 * nrow(Xnew)) {
        linpred <- function(p) (Xnew + SEnew * Crnorm(length(SEnew))) %m*v% p[[name]]
        linpred_update <- function(x, plus=TRUE, p) mv_update(x, plus, Xnew + SEnew * Crnorm(length(SEnew)), p[[name]])
      } else {
        i.me.new <- whichv(dapply(SEnew, \(x) all(x > 0), MARGIN=1L), TRUE)  # TODO add tolerance + allow list of matrix
        if (length(i.me.new) == nrow(newdata)) i.me.new <- seq_len(nrow(newdata))
        SEnew <- SEnew[i.me.new, ]
        linpred <- function(p) {
          out <- Xnew
          out[i.me.new, ] <- out[i.me.new, ] + SEnew * Crnorm(length(SEnew))
          out %m*v% p[[name]]
        }
        linpred_update <- function(x, plus=TRUE, p) {
          out <- Xnew
          out[i.me.new, ] <- out[i.me.new, ] + SEnew * Crnorm(length(SEnew))
          mv_update(x, plus, out, p[[name]])
        }
      }
    }
    rm(newdata, verbose)
    environment()
  }

  rprior <- function(p) {}
  # X, QME are assumed to be part of the prior information
  if (nme == fam[["n"]]) {
    rprior <- add(rprior, bquote(p[[.(name_X)]] <- X + sqrt(.(if (q == 1L) quote(1/QME) else quote(V))) * Crnorm(.(nme*q))))
  } else {
    rprior <- rprior |>
      add(bquote(p[[.(name_X)]] <- X)) |>
      add(bquote(p[[.(name_X)]][i.me, ] <- X[i.me, ] + sqrt(.(if (q == 1L) quote(1/QME) else quote(V))) * Crnorm(.(nme*q))))
  }
  rprior <- rprior |>
    add(bquote(p[[.(name)]] <- prior$rprior(p))) |>
    add(quote(p))

  if (!is.null(constraints)) {
    # R, S, r, s used in mc_block.R
    if (constraints[["eq"]]) {
      R <- constraints[["R"]]
      r <- constraints[["r"]]
    }
    if (constraints[["ineq"]]) {
      S <- constraints[["S"]]
      s <- constraints[["s"]]
    }
  }

  self <- environment()
  if (prior.only) return(self)

  if (in.block && use.sigma)
    get_Q <- function(p) Q0@x * (1/p[["sigma_"]]^2)

  draw <- if (debug) function(p) {browser()} else function(p) {}
  if (use.sigma && prior[["type"]] == "normal" && informative.prior) {
    # in case of a gaussian model the variance of normal priors is taken
    # proportional to sigma^2 --> define contributions to sigma posterior
    if (!is.null(constraints) && constraints[["eq"]])
      fam$df.sigma <- fam[["df.sigma"]] + q - constraints[["ncR"]]
    else
      fam$df.sigma <- fam[["df.sigma"]] + q
    SSR_sigma <- function(p) {
      delta.beta <- if (zero.mean) p[[name]] else p[[name]] - prior[["mean"]]
      dotprodC(delta.beta, Q0 %m*v% delta.beta)
    }
  }

  if (sc[["single.block"]] && sc[["length1mod"]]) {
    if (fam[["e.is.res"]]) draw <- add(draw, quote(p$e_ <- copy_obj(fam[["y"]])))
  } else {
    if (fam[["e.is.res"]])
      draw <- add(draw, bquote(mv_update(p[["e_"]], plus=TRUE, p[[.(name_X)]], p[[.(name)]])))
    else
      draw <- add(draw, bquote(mv_update(p[["e_"]], plus=FALSE, p[[.(name_X)]], p[[.(name)]])))
  }

  # 1. draw covariates p[[name_X]]
  # for now we assume that X is dense
  if (!fam[["e.is.res"]] && sc[["single.block"]] && !sc[["length1mod"]]) {
    # need to correct Q_e function; cf gen component with PX
    if (fam[["link"]] == "probit")
      draw <- add(draw, quote(Qe <- fam$Q_e(p) - p[["e_"]]))
    else
      draw <- add(draw, quote(Qe <- fam$Q_e(p) - p[["Q_"]] * p[["e_"]]))
  } else {
    draw <- add(draw, quote(Qe <- fam$Q_e(p)))
  }
  if (fam[["modeled.Q"]])
    draw <- add(draw, quote(v0 <- 1/p[["Q_"]][i.me]))
  draw <- add(draw, bquote(scale <- .(if (fam[["sigma.fixed"]]) quote(v0) else quote(v0 * p[["sigma_"]]^2))))
  if (q == 1L) {
    draw <- add(draw, bquote(Vscaled <- 1 / (QME * scale + p[[.(name)]]^2)))
    if (nme == fam[["n"]])
      draw <- add(draw, bquote(p[[.(name_X)]] <- X + Vscaled * p[[.(name)]] * (v0 * Qe - p[[.(name)]] * X) + sqrt(scale * Vscaled) * Crnorm(.(nme))))
    else
      draw <- add(draw, bquote(p[[.(name_X)]][i.me] <- X[i.me] + Vscaled * p[[.(name)]] * (v0 * Qe[i.me] - p[[.(name)]] * X[i.me]) + sqrt(scale * Vscaled) * Crnorm(.(nme))))
  } else {
    # V-form (independent m.e.)
    draw <- add(draw, bquote(
      for (i in i.me) {
        Vi <- V[i, ]
        Cb <- Vi * p[[.(name)]]
        f <- 1 / (scale[i] + dotprodC(p[[.(name)]], Cb))
        Vx <- add_diagC(- f * base_tcrossprod(Cb), Vi)
        Xi <- X[i, ]
        p[[.(name_X)]][i, ] <- Xi + Cb * f * (p[["e_"]][i] - dotprodC(p[[.(name)]], Xi)) + crossprod_mv(chol.default(Vx), rnorm(q))
      }
    ))
    # TODO use Rcpp, and use Q-form at least for cases where all QME[i, ] > 0, and handle list-of-matrix (correlated m.e.)
  }

  if (!in.block) {
    # 2. draw coefficients p[[name]]
    if (allv(Q0b0, 0))
      draw <- add(draw, bquote(Xy <- crossprod_mv(p[[.(name_X)]], Qe)))
    else
      draw <- add(draw, bquote(Xy <- crossprod_mv(p[[.(name_X)]], Qe) + Q0b0))

    if (fam[["modeled.Q"]]) {
      if (fam[["Q0.type"]] == "symm")
        draw <- add(draw, bquote(XX <- crossprod_sym(p[[.(name_X)]], p[["QM_"]])))
      else
        draw <- add(draw, bquote(XX <- crossprod_sym(p[[.(name_X)]], p[["Q_"]])))
    } else {
      draw <- add(draw, bquote(XX <- crossprod_sym(p[[.(name_X)]], fam[["Q0"]])))
    }
    if (informative.prior) {
      # TODO instead of runif(n, ...) account for the structure in Qmod; lambda may not vary per unit!
      # TODO here we need M1 to be dense without zeros
      mat_sum <- make_mat_sum(M0=Q0, M1=crossprod_sym(X, crossprod_sym(Cdiag(runif(fam[["n"]], 0.9, 1.1)), fam[["Q0"]])))
      MVNsampler <- create_TMVN_sampler(
        Q=mat_sum(crossprod_sym(X, crossprod_sym(Cdiag(runif(fam[["n"]], 0.9, 1.1)), fam[["Q0"]]))),
        update.Q=TRUE, name=name, constraints=constraints
      )
      draw <- add(draw, quote(MVNsampler$update(mat_sum(XX))))
    } else {
      MVNsampler <- create_TMVN_sampler(
        Q=crossprod_sym(X, crossprod_sym(Cdiag(runif(fam[["n"]], 0.9, 1.1)), fam[["Q0"]])),
        update.Q=TRUE, name=name, constraints=constraints
      )
      draw <- add(draw, quote(MVNsampler$update(XX)))
    }
    draw <- add(draw, bquote(p[[.(name)]] <- MVNsampler$draw(p, .(if (fam[["sigma.fixed"]]) 1 else quote(p[["sigma_"]])), Xy=Xy)[[.(name)]]))
  }  # END if (!in.block)

  # compute full residuals/fitted values
  if (fam[["e.is.res"]])
    draw <- add(draw, bquote(mv_update(p[["e_"]], plus=FALSE, p[[.(name_X)]], p[[.(name)]])))
  else if (sc[["single.block"]] && sc[["length1mod"]])
    draw <- add(draw, bquote(p[["e_"]] <- p[[.(name_X)]] %m*v% p[[.(name)]]))
  else
    draw <- add(draw, bquote(mv_update(p[["e_"]], plus=TRUE, p[[.(name_X)]], p[[.(name)]])))
  draw <- add(draw, quote(p))

  start <- function(p) {
    if (!in.block) p <- MVNsampler$start(p)
    p[[name_X]] <- X
    p
  }

  self
}