R/gAnalyticSolve.R

Defines functions gAnalyticSolve devG

Documented in devG gAnalyticSolve

#' The main function which calculates the analytic solution to the linear mixed effects model. 
#' @param weights level-1 weights
#' @param y outcome measure. 
#' @param X the X matrix.
#' @param Zlist, a list of matrixes with the Z values for each level. 
#' @param Zlevels, the level corresponding to each matrix in Zlist. 
#' @param weights a list of unconditional weights for each model level. 
#' @param weightsC a list of conditional weights for each model level. 
#' @param  groupID a matrix containing the group ids for each level in the model. 
#' @param  lmeVardf a dataframe containing the variances and covariance of the random effects, in the same format as returned from lme. 
#' @importFrom Matrix Diagonal Matrix
#' @importFrom methods as
#' @keywords internal
gAnalyticSolve <- function(y, X, Zlist, Zlevels, weights, weightsC=weights, pWeights=NULL, pred0, family, groupID, lmeVarDF) {
  
  if(!inherits(groupID, "matrix")) {
    stop(paste0("Variable", dQuote(groupID)," must be a matrix with IDs at a level per column."))
  }
  if(!inherits(y, "numeric")) {
    stop(paste0(dQuote("y"), " must be a numeric vector."))
  }
  if(!is.null(pWeights)) {
    weights[[1]] <- weights[[1]] * pWeights
    weightsC[[1]] <- weightsC[[1]] * pWeights
  }

  ny <- length(y)
  X <- as.matrix(X)
  nX <- nrow(X)
  if(nX != ny) {
    stop(paste0("Length of the ", dQuote("y"), " vector and the ", dQuote("X"), " matrix must agree."))
  }
  if(length(Zlist) != length(Zlevels)) {
    stop(paste0("Length of the ", dQuote("Zlist"), " and ", dQuote("Zlevels"), " must agree."))
  }
  nW1 <- length(weights[[1]])
  if(nW1 != nX) {
    stop(paste0("Number of rows in ", dQuote("weights[[1]]") , " must agree with number of rows in ", dQuote("X"),"."))
  }
  nWc1 <- length(weightsC[[1]])
  if(nWc1 != nX) {
    stop(paste0("Number of rows in ", dQuote("weightsC[[1]]") , " must agree with number of rows in ", dQuote("X"),"."))
  }
  ngid <- nrow(groupID)
  if(ngid != nX) {
    stop(paste0("Number of rows in ", dQuote("groupID") , " must agree with number of rows in ", dQuote("X"),"."))
  }
  
  # get the number of groups per level
  nc <- apply(groupID, 2, function(x) {length(unique(x))} )
  
  # for level 1 cast Z for lowest level as a matrix and transpose to get Zt
  Zt <- Matrix::t(as(Zlist[[1]], "Matrix"))
  
  # for level, build PsiVec, the diagonal of the Psi  (weights) matrix
  PsiVec <- rep(weights[[Zlevels[1]]], ncol(Zlist[[1]])/nc[1])
  
  # Assemble the Zt matrix and psi Vector for levels >1 
  ZColLevels <- rep(Zlevels[1], ncol(Zlist[[1]]))
  if(length(Zlist) > 1) {
    for(zi in 2:length(Zlist)) {
      Zt <- rbind(Zt, Matrix::t(as(Zlist[[zi]], "Matrix")))
      ZColLevels <- c(ZColLevels, rep(Zlevels[zi], ncol(Zlist[[zi]])))
      PsiVec <- c(PsiVec, rep(weights[[Zlevels[zi]]], ncol(Zlist[[zi]])/nc[Zlevels[zi]-1]))
    }
  }
 
  # get unit weights
  W0 <- weights[[1]]
  Psi <- Diagonal(n=nrow(Zt), x=PsiVec) #matrix of weights at level one
  Psi12 <- Diagonal(n=nrow(Zt), x=sqrt(PsiVec)) #matrix of square root of level one weights
  W <- Diagonal(x=(W0))
  W12 <- Diagonal(x=sqrt(W0))
  
  # calculate  conditional weights matrix where level one weights are scaled by level two wieghts
  W12cDiag <- W0
  L1groups <- unique(groupID[,1])
  for(gi in 1:length(L1groups)) {
    rowsGi <- groupID[,1] == L1groups[gi]
    W12cDiag[rowsGi] <- sqrt(W12cDiag[rowsGi] / weights[[2]][gi])
  }
  W12c <- Diagonal(x=W12cDiag)
  Z <- Matrix::t(Zt)
  options(Matrix.quiet.qr.R=TRUE) #supress extra print outs 

  M0 <- Matrix(data=0, ncol=ncol(X), nrow=nrow(Zt))
  
  #discrepency function for calculating least squares solution. Follows from Bates et al. 2015 eq 14 and 15
  discf <- function(Zt, X, lambda, u, Psi12, W12, b) { # return ehatT * ehat
    wres <- W12 %*% (y - Matrix::t(Zt) %*% lambda %*% u - X %*% b) # residual
    ures <- Psi12 %*% u # augmented ehat
    as.numeric(sum(wres^2) + sum(ures^2))
  }
  
  function(v, verbose=FALSE, beta=NULL, sigma=NULL, robustSE=FALSE) {
    beta0 <- beta
    sigma0 <- sigma
    omegaVec <- v

    #Create list delta and lambda Matrixes for each level
    Delta <- list()
    lambda_by_level <- list()
    
    levels <- max(lmeVarDF$level)
   
    for (l in 2:levels){
      #set up matrix of zeros with rows and columns for each random effect
      #number of random effect is the number of variance terms at that level (not including covariance terms )
      n_ref_lev  <- nrow(lmeVarDF[lmeVarDF$level==l & is.na(lmeVarDF$var2),] ) 
      lambda_i  <- matrix(0,nrow=n_ref_lev,ncol=n_ref_lev)
      row.names(lambda_i) <- lmeVarDF[lmeVarDF$level==l & is.na(lmeVarDF$var2),"var1"]
      colnames(lambda_i) <- lmeVarDF[lmeVarDF$level==l & is.na(lmeVarDF$var2),"var1"]
      
      #get only v for this level
      group <- lmeVarDF[lmeVarDF$level==l ,"grp"][1]
      v_lev <- v[grep(paste0("^",group,"."),names(v))]
      
      
      #fill in lambda_i from theta using names
      for (vi in 1:length(v_lev)){
        row_index <- strsplit(names(v_lev[vi]),".",fixed=TRUE)[[1]][2]
        col_index <- ifelse(length(strsplit(names(v_lev[vi]),".",fixed=TRUE)[[1]])>2,strsplit(names(v_lev[vi]),".",fixed=TRUE)[[1]][3],row_index )
        lambda_i[row_index,col_index] <- v_lev[vi]
      }
      diag(lambda_i)[diag(lambda_i)==0] <- .Machine$double.eps #set any 0s to smallest non zero number to enable solve in next step
      Delta[[l]] <- solve(lambda_i)
      
      #assemble blockwise lambad by level matrix 
      lambda_by_level[[l]] <- kronecker(lambda_i,diag(1,unique(lmeVarDF[lmeVarDF$level==l & is.na(lmeVarDF$var2),"ngrp"])))
    }
  
    #Create lambda matrix from diagonal of all lamdas 
    lambda <- bdiag(lambda_by_level[2:length(lambda_by_level)]) # vignette equation 21

    etaHat <- family$linkfun(pred0)
    yprime <- etaHat + (y - pred0) * ll$mu.eta(etaHat)

    yaugmented <- c(as.numeric(W12 %*% yprime), rep(0, nrow(Zt))) # first matrix in vingette equation 25 and 60 
    A <- rbind(cbind(W12 %*% Z %*% lambda, W12 %*% X),
               cbind(Psi12,M0)) # second matrix in equation (60)/(25). 
    qr_a <- qr(A)

    
    ub <- qr.coef(qr_a, yaugmented) # equation (26)/ stack vector (u: top part for fixed effects coef estimate, b: bottom part for random effects vector (i.e. of length n_group * n_random_effects)
    return(ub[(1:ncol(Z))])
    # get ln of determinant of Lz matrix (Bates et al. 2015 notation) / calculation of alpha
    # or the final product in Pinheiro and Bates 2.13 (formula used for weights)
    lndetLz <- 0 # lndetLz is alpha in the specs
    lndetLzg <- list() # by top level group lnldetLz
    # used to find columns with no-non-zero elements
    nozero <- function(x) { any(x != 0) }
    for(level in 1:ncol(groupID)) {
      groups <- unique(groupID[,level])
      Deltai <- Delta[[level+1]]
      # R11 notation from Bates and Pinheiro, 1998
      R11 <- list()
      topLevel <- level == ncol(groupID)
      for(gi in 1:length(groups)) {
        # this works for 3 level, will need to consider selection by row for >3 
        Zrows <- groupID[,level]==groups[gi]
        # filter to the rows for this group, also filter to just columns
        # for this level of the model
        Zi <- Z[Zrows,ZColLevels==(level+1),drop=FALSE] # Zi associated with the random effect (~ Z_g in the specs)
        # within columns for this level, only the non-zero columns
        # regard this unit, so filter it thusly
        Zcols <- apply(Zi,2,nozero)
        Zi <- Zi[,Zcols,drop=FALSE]

        if(level == 1 || level < ncol(groupID)) {
          if(!topLevel) {
            lp1 <- unique(groupID[Zrows,level+1])
            if(length(lp1) > 1) {
              stop("Not a nested model.")
            }
            qp1 <- nrow(Delta[[level+2]])
          }
          qi <- nrow(Deltai) 
          Zi <- Z[Zrows,,drop=FALSE]
          # within columns for this level, only the non-zero columns
          # regard this unit, so filter it thusly
          Zcols <- apply(Zi,2,nozero)
          Zi <- Zi[,Zcols,drop=FALSE]
          # update W^{1/2}, for this purpose, it's the conditional weights
          W12i <- W12c[Zrows,Zrows]
          ZiA <- W12i %*% Zi
          #shape delta 
          Delta0 <- Matrix(data=0, nrow=qi,ncol=ncol(ZiA)-ncol(Deltai))
          ZiA <- rbind(ZiA, cbind(Deltai, Delta0))
  
          #in this section we decompose Zi * A using qr decomposition, following equation 43 in the vingette. 
          ZiAR <- qr.R(qr(ZiA))
          R22 <- ZiAR[1:qi, 1:qi, drop=FALSE]
          lndetLzi <- weights[[level+1]][gi]*(- as.numeric(determinant(Deltai)$mod) + sum(log(abs(diag(R22)[1:ncol(Deltai)])))) # equation (92)
          # save R11 for next level up, if there is a next level up
          if(!topLevel) {
            R11i <- sqrt(weightsC[[level+1]][gi])*ZiAR[-(1:qi),qi+1:qp1, drop=FALSE]
            # load in R11 to the correct level
            if(length(R11) >= lp1) {
              R11[[lp1]] <- rbind(R11[[lp1]], R11i)
              lndetLzg[[lp1]] <- lndetLzg[[lp1]] + lndetLzi
            } else {
              # there was no R11, so start it off with this R11
              R11[[lp1]] <- R11i
              lndetLzg[[lp1]] <- lndetLzi
            }
          } else {
            lndetLzg[[gi]] <- lndetLzi
          }
        }
        if(level>=2) {
          ZiA <- rbind(pR11[[groups[gi]]], Deltai)
          R <- qr.R(qr(ZiA))
          # weight the results
          lndetLzi <- weights[[level+1]][gi]*(- (as.numeric(determinant(Deltai)$mod)) + sum(log(abs(diag(R)[1:ncol(Deltai)])))) # equation 92
          lndetLzg[[gi]] <- lndetLzg[[gi]] + lndetLzi
        }
        # update W^{1/2}, for this purpose, it's the conditional weights
        # W12i <- W12c[Zrows,Zrows]
        lndetLz <- lndetLz + lndetLzi
        if(verbose) {
          cat("gi=", gi, " lndetLzi=", lndetLzi, " lndetLz=", lndetLz,"w=", weights[[level+1]][gi], "\n")
        }
      }
      if(level < ncol(groupID)) {
        pR11 <- R11
      }
    }
    # this is the beta vector
    b <- ub[-(1:ncol(Z))]
    if(!is.null(beta)) {
      if(length(b) != length(beta)) {
        stop(paste0("The argument ", dQuote("beta"), " must be a vector of length ", length(b), "."))
      }
    }
    # and the random effect vector
    u <- ub[1:ncol(Z)]
  
    
    # wrap the random effect vector to matrix format so that each row regards one group
    tmpu <- u
    bb <- list()
    vc <- matrix(nrow=1+length(v), ncol=1+length(v))
    u0 <- 0
    for(li in 2:length(lambda_by_level)) {
      lli <- lambda_by_level[[li]]
      ni <- nrow(lli)
      bb[[li]] <- lli %*% u[u0 + 1:ni]
      u0 <- u0 + ni
      vc[li,li] <- 1/(ni) * sum((bb[[li]])^2) # mean already 0
    }

    # the discrepancy ||W12(y-Xb-Zu)||^2 + ||Psi(u)||^2
    discrep <- discf(Zt, X, lambda, u, Psi12, W12, b)
    if(!is.null(beta) & is.null(sigma)) {
      discrep <- discrep + sum( (RX %*% (beta - b))^2 )
    }
    # the R22 matrix, bottom right of the big R, conforms with b
    RX <- qr.R(qr_a)[(ncol(Z)+1):ncol(A),(ncol(Z)+1):ncol(A)]
    nx <- sum(W0)
    
    # residual
    sigma <- sqrt(discrep / nx)
    dev <- 0
    if(!is.null(sigma0)) {
      sigma <- sigma0
      dev <- 2*as.numeric(lndetLz) + nx * log(2*pi*(sigma^2)) + discrep/sigma^2 + sum((RX %*% (beta - b))^2)/sigma^2  # equation (98d)
      if(verbose) {
        cat("lnl2=", (nx * log(2*pi*(sigma^2)) + discrep/sigma + sum((RX %*% (beta - b))^2)/sigma)/-2,"\n")
      }
    } else {
      dev <- 2*as.numeric(lndetLz) + nx * (log(2*pi*(sigma^2)) + 1)  # equation (98)
      if(verbose) {
        cat("lnl2=", (nx * (log(2*pi*(sigma^2)) + 1)/-2),"\n")
      }
    }

    
    #Variance of Beta, from Bates et.al. 2015
    cov_mat <- (sigma^2) * solve(RX) %*% solve(t(RX))
    varBeta <- diag(cov_mat)

    sigma <- sqrt(discrep / nx)

    res <- list(b=b, u=u, ranef=bb, discrep=discrep, sigma=sigma, dev=dev,
                lnl=dev/-2, theta=v, varBeta=varBeta, vars=NULL, cov_mat=cov_mat,
                lndetLz=lndetLz)
    if(robustSE) {
      #Calculate robust standardize effect
      # based on the sandwich estimator of Rabe-Hesketh and Skrondal 2006
      # this whole calculation focuses on calculating the liklihood at the top group level 
      fgroupID <- groupID[,ncol(groupID)] 
      uf <- unique(fgroupID) # unique final groupIDs
      lnli <- vector(length=length(uf))
      Jacobian <- matrix(NA, nrow=length(uf), ncol=length(b))
      
      #this code block seperates wieghts into the set of weights belonging to each top level group 
      for(gi in 1:length(uf)) {
        sgi <- fgroupID == gi # subset for group gi
        weightsPrime <- list(weights[[1]][sgi])
        weightsCPrime <- list(weightsC[[1]][sgi])
        for(i in 2:length(weights)) {
          theseGroups <- unique(groupID[sgi,i-1])
          weightsPrime[[i]] <- weights[[i]][theseGroups]
          weightsCPrime[[i]] <- weightsC[[i]][theseGroups]
        }
        condenseZ <- function(z) {
          res <- z[sgi,]
          res[,apply(res,2,sum)>0, drop=FALSE]
        }
        groupIDi <- groupID[sgi,,drop=FALSE]
        lmeVarDFi <- lmeVarDF
        lmeVarDFi$ngrp[lmeVarDFi$level==1] <- sum(sgi)
        for(i in 2:length(weights)) {
          lmeVarDFi$ngrp[lmeVarDFi$level==i] <- length(unique(groupIDi[,i-1]))
        }
        #calculate the group level likelihood by applying the function to only the x or y within the group indexed by sgi
        bwi <- analyticSolve(y=y[sgi], X[sgi,],
                  Zlist=lapply(Zlist, condenseZ),
                  Zlevels=Zlevels,
                  weights=weightsPrime,
                  weightsC=weightsCPrime,
                  groupID=groupIDi,
                  lmeVarDF=lmeVarDFi) 
        lnli[gi] <- bwi(v=v, verbose=verbose,
                        beta=b, sigma=sigma,
                        robustSE=FALSE)$lnl
        bwiW <- function(bwi, v, sigma, b, ind) {
          function(par) {
            b[ind] <- par
            bwi(v=v, verbose=FALSE,
                b=b,
                sigma=sigma,robustSE=FALSE)$lnl
          }
        }
        for(j in 1:length(b)){
          Jacobian[gi,j] <- d(bwiW(bwi, v=v, sigma=sigma, b=b, ind=j), par=b[j])
        }
      }
      J <- matrix(0,ncol=length(b), nrow=length(b))
      for (i in 1:nrow(Jacobian)){
        J <- J + Jacobian[i,] %*% t(Jacobian[i,])
      }
      J <- (nrow(Jacobian)/(nrow(Jacobian)-1))*J
      varBetaRobust <- cov_mat %*% J %*% cov_mat
      res <- c(res, list(varBetaRobust=varBetaRobust, seBetaRobust=sqrt(diag(varBetaRobust))))
    }
    return(res)
  }
}

#' a wrapper for analyticSolve that allows it to be called from an optimizer. Takes the same arguments as analyticSolve. 
#' @param weights level-1 weights
#' @param y outcome measure. 
#' @param X the X matrix.
#' @param Zlist, a list of matrixes with the Z values for each level. 
#' @param Zlevels, the level corresponding to each matrix in Zlist. 
#' @param weights a list of unconditional weights for each model level. 
#' @param weightsC a list of conditional weights for each model level. 
#' @param  groupID a matrix containing the group ids for each level in the model. 
#' @param  lmeVardf a dataframe containing the variances and covariance of the random effects, in the same format as returned from lme. 
#' @importFrom Matrix Diagonal Matrix
#' @importFrom methods as
#' @keywords internal
devG <- function(y, X, Zlist, Zlevels, weights, weightsC=weights, groupID,lmeVarDF) {
  bs <- analyticSolve(y=y, X=X, Zlist=Zlist, Zlevels=Zlevels, weights=weights, weightsC=weightsC, lmeVarDF=lmeVarDF,
           groupID=groupID)
  function(v, getBS=FALSE) {
    if(getBS) {
      return(bs)
    }
    bhat <- bs(v)
    return(bhat$dev)
  }
}
ClaireKelley/WeMix documentation built on May 21, 2019, 6:46 a.m.