R/fitBLUP.R

In MarcooLopez/SFSI_data: Sparse Family and Selection Index

# X= Z = U = d = indexK = NULL; BLUP=FALSE; method="ML"; h2=NULL
# return.Hinv = FALSE; tol=1E-5; maxIter=1000; interval=c(1E-9,1E9)

fitBLUP <- function(y, X = NULL, Z = NULL, K = NULL, U = NULL, d = NULL,
                        indexK = NULL, h2 = NULL, BLUP = TRUE, method = "ML",
                        return.Hinv = FALSE, tol = 1E-5, maxIter = 1000,
                        interval = c(1E-9,1E9), warn = TRUE)
{
  method <- match.arg(method)
  eps <- .Machine$double.eps

  if(is.character(K)){
    K <- readBinary(K,indexRow=indexK,indexCol=indexK)
  }
  y <- as.vector(y)
  indexNA <- which(is.na(y))
  indexOK <- which(!is.na(y))
  n <- length(indexOK)

  if(is.null(X))
  {
    X <- stats::model.matrix(~1,data=data.frame(rep(1,length(y))))
  }else{
    if(is.vector(X)){
      X <- stats::model.matrix(~X)
      if(ncol(X)>2)  colnames(X)[-1] <- substr(colnames(X)[-1],2,nchar(colnames(X)[-1]))
    }
  }
  stopifnot(nrow(X) == length(y))

  if(is.null(U) & is.null(d))
  {
    if(is.null(Z))
    {
      Z <- diag(length(y))
      if(is.null(K)){
          K <- diag(length(y))
      }
      G <- K
    }else{
      if(!is.matrix(Z)) stop("Object 'Z' must be a matrix")
      if(is.null(K)){
        K <- diag(ncol(Z))
        G <- tcrossprod(Z)  # K = ZIZ'
      }else{
        G <- Z%*%K%*%t(Z)
      }
    }
    stopifnot(nrow(G) == length(y))
    stopifnot(ncol(G) == length(y))

    out <- eigen(G[indexOK,indexOK])
    d <- out$values
    U <- out$vectors

  }else{
    if(is.null(Z))  Z <- diag(length(y))
    if(is.null(U)) stop("You are providing the eigenvalues, but not the eigenvectors")
    if(is.null(d)) stop("You are providing the eigenvectors, but not the eigenvalues")
    if(length(indexNA)>0) stop("No 'NA' values are allowed when passing parameters 'U' and 'd'")
  }

  stopifnot(nrow(U) == n)
  stopifnot(ncol(U) == length(d))

  Uty <- drop(crossprod(U,y[indexOK]))
  UtX <- crossprod(U,X[indexOK, ,drop=FALSE])

  # Function to search in a given subintervals
  searchInt <- function(bb){
    flag <- TRUE; i <- 1
    convergence <- lambda0 <- dbar <- ytPy <- bHat <- NA
    while(flag)
    {
      i <- i + 1
      tmp <- try(uniroot(f=dloglik,interval=c(bb[i-1],bb[i]),n=n,Uty=Uty,
                         UtX=UtX,d=d,tol=tol,maxiter=maxIter,trace=2),
                 silent = TRUE)
      if(class(tmp) == "list")
      {
        lambda00 <- tmp$root
        if(lambda00 <= interval[1]){
          lambda00 <- interval[1]
        }else{
          if(lambda00 >= interval[2]){
            lambda00 <- interval[2]
          }
        }

        dbar <- 1/(lambda00*d + 1)
        qq1 <- t(Uty*dbar)%*%UtX
        qq2 <- solve(sweep(t(UtX),2L,dbar,FUN="*")%*%UtX)
        ytPy <- drop(sum(dbar*Uty^2)-qq1%*%qq2%*%t(qq1))
        bHat <- drop(qq2%*%t(qq1))

        varP <- var(y[indexOK])*mean(d)
        if(lambda00*ytPy/n <= (1.05)*varP){
          convergence <- tmp$iter <= maxIter
          lambda0 <-  lambda00
        }
      }
      #aa <- rep(NA,3)
      #if(class(tmp) == "list") aa=c(tmp$root,tmp$f.root,tmp$estim.prec)
      #cat("Interval ",i-1,"[",bb[i-1],",",bb[i],"]: root=",aa[1]," f.root=",aa[2]," prec=",aa[3],"\n")
      if(i == length(bb) | !is.na(convergence)) flag <- FALSE
    }
    list(lambda0=lambda0,convergence=convergence,dbar=dbar,ytPy=ytPy,bHat=bHat)
  }

  convergence <- lambda0 <- ytPy <- bHat <- dbar <- NA
  if(is.null(h2))
  {
    tt <- searchInt(interval)
    convergence <- tt$convergence

    if(is.na(convergence))
    {
      # Divide seeking interval into smaller intervals
      bb <- exp(seq(log(interval[1]),log(interval[2]),length=200))
      tt <- searchInt(bb)
      convergence <- tt$convergence

      if(is.na(convergence)){
        # Search in the lower bound
        bb <- exp(seq(log(interval[1]^2),log(interval[1]),length=200))
        tt <- searchInt(bb)
        convergence <- tt$convergence

        if(is.na(convergence)){
          # Search in the upper bound
          bb <- exp(seq(log(interval[2]),log(interval[2]^2),length=200))
          tt <- searchInt(bb)
          convergence <- tt$convergence
        }
      }
    }
    lambda0 <- tt$lambda0
    ytPy <- tt$ytPy
    bHat <- tt$bHat
    dbar <- tt$dbar

  }else{
    lambda0 <- h2/(1-h2)
    dbar <- 1/(lambda0*d + 1)
    qq1 <- t(Uty*dbar)%*%UtX
    qq2 <- solve(sweep(t(UtX),2L,dbar,FUN="*")%*%UtX)
    ytPy <- drop(sum(dbar*Uty^2)-qq1%*%qq2%*%t(qq1))
    bHat <- drop(qq2%*%t(qq1))
  }

  # Compute BLUP: uHat = KZ' V^{-1} (y-X*b)
  uHat <- yHat <- Hinv <- NULL
  if(BLUP & !is.na(lambda0))
  {
    yStar <- y[indexOK] - X[indexOK ,,drop=FALSE]%*%bHat
    Hinv <- tcrossprod(sweep(U,2L,lambda0*dbar,FUN="*"),U) # Vinv
    #H <- tcrossprod(sweep(U,2L,d*lambda0*dbar,FUN="*"),U)

    if(is.null(K)) K <- tcrossprod(sweep(U,2L,d,FUN="*"),U)
    KZt <- tcrossprod(K,Z[indexOK,,drop=FALSE])

    uHat <- drop(KZt%*%Hinv%*%yStar)   # u = K Z' Vinv y*
    yHat[indexOK] <- drop(X[indexOK, ,drop=FALSE]%*%bHat + Z[indexOK, ,drop=FALSE]%*%uHat)
    if(length(indexNA)>0){
      yHat[indexNA] <- drop(X[indexNA,,drop=FALSE]%*%bHat + Z[indexNA, ,drop=FALSE]%*%uHat)
    }

    if(!return.Hinv) Hinv <- NULL
  }

  if(warn){
    if(ifelse(is.na(convergence),is.na(lambda0),!convergence)){
      warning("Convergence was not reached in the 'EMMA' algorithm.",immediate.=TRUE)
     }
  }

  varE <- ytPy/n
  varU <- lambda0*varE
  h2 <- varU/(varU + varE)

  out <- list(varE = varE, varU = varU, h2 = h2, yHat=yHat, b = bHat, u = uHat,
              Hinv = Hinv, convergence = convergence, method = method)
  return(out)
}