R/fit_LMM.R

In devoges/fstat: fstat

#' Fit linear mixed model
#'
#' @param Y n by 1 vector
#' @param X n by p matrix
#' @param Zs list of n by i_k matrices. NAMED
#' @param Ss list of n by n matrices, NAMED
#' @param prefix what to append to beginning of messages
#' @param print whether to print information
#'
#' @return
#' @export
fit_LMM <- function(Y,
                    X,
                    Zs = NULL,
                    Ss = NULL,
                    prefix = NULL,
                    print = FALSE) {
  out <- list()

  if ( length(Zs) > 0 & length(Ss) == 0 ) {

    if ( print ) {
      message( prefix %% ": using lme4")
    }
    model <- fit_LMER(Response = Y,
                      X = X,
                      ListZ = Zs)

    out$method <- "lme4"
    out$raw_model <- model
    out$beta <- getME(model, name = "fixef")
    out$vcs <- as.data.frame(lme4::VarCorr(model))[, "vcov"] %>%
      `names<-`(c(remove_last(as.data.frame(lme4::VarCorr(model))[, "var1"]), "Error"))
    out$RSS <- unname(c(lme4::getME(model, "devcomp")$cmp["pwrss"]))
    out$rml <- logLik(model, REML = TRUE)[1]
    out$ml <- logLik(model, REML = FALSE)[1]

  } else if ( length(Ss) > 0 ) {
    if ( print ) {
      message(prefix %% ": using GMMAT")
    }

    if ( length(Zs) > 0 ) {
      message(prefix %% ": converting Zs to Ss")

      Ss_from_Zs <- nlapply(names(Zs), function(x) {
        tcrossprod(Zs[[x]])
      })

      Ss <- c(Ss, Ss_from_Zs)
    }

    model <- suppressWarnings(fit_GMMAT(Y = Y,
                                        X = X,
                                        Ss = Ss))

    out$method <- "GMMAT"
    out$raw_model <- model
    out$beta <- model$coefficients %>% `names<-`(colnames(X))
    out$vcs <- model$theta
    out$RSS <- model$RSS
    out$rml <- model$rml
    out$ml <- model$ml
  } else if ( length(Zs) == 0 & length(Zs) == 0 ) {
    if ( print ) {
      message( prefix %% ": using lm")
    }

    model <- lm(Y ~ X - 1)

    out$method <- "lm"
    out$raw_model <- model
    out$beta <- coef(model) %>% `names<-`(colnames(X))
    out$vcs <- summary(model)$sigma %>% `names<-`("Error")
    out$RSS <- sum(resid(model)^2)

    ##-- Calculate RMV
    V <- diag(nrow(X))
    Vinv <- minv(V)
    P <- diag(nrow(X)) - X %mtimes% minv(t(X) %mtimes% Vinv %mtimes% X) %mtimes% t(X) %mtimes% Vinv
    detV <- determinant(V, logarithm = TRUE)$modulus
    logYPVPY <- log(t(Y) %mtimes% t(P) %mtimes% Vinv %mtimes% P %mtimes% Y)
    rml <- -1 * detV +
      -1 * determinant(t(X) %mtimes% Vinv %mtimes% X, logarithm = TRUE)$modulus +
      -1 * (nrow(X) - ncol(X)) * logYPVPY


    out$rml <- logLik(model, REML = TRUE)[1]
    out$rml <- rml[1, 1]/2
    out$ml <- logLik(model, REML = FALSE)[1]
  }

  return(out)
}

# debugonce(fit_LMM)
# source("R/fixGMMAT.R")
# source("R/fixlmer.R")
# source("R/simData.R")
#
# dvRutils::library_many(c("dvRutils",
#                          "MASS",
#                          "Matrix",
#                          "magrittr",
#                          "lme4",
#                          "tibble",
#                          "GMMAT"))
# # set.seed(21242)
# d <- sim_data(n = 200,
#               p = 2,
#               beta = c(1, 5, 1),
#               Z_ks = c(1, 1),
#               Z_sigmas = c(0.5, 0.5),
#               S_n = 1,
#               S_sigmas = c(1),
#               sigE = 1)
# Y <- d$Y
# X <- d$X
# Zs <- d$Zs
# Ss <- d$Ss
# m <- fit_LMM(Y = Y,
#              X = X,
#              Ss = Ss,
#              Zs = Zs,
#              print = TRUE)