R/impute.R

In chr1swallace/annotSnpStats: Annotated SnpMatrix objects

##' Uses snpStats' snp.imputation function
##'
##' @title Fill in missing values in aSnpMatrix object
##' @param X SnpMatrix or aSnpMatrix object
##' @param bp optional vector giving bp of SNPs, of length == number of SNPs
##' @param strata optional strata vector for imputation within strata, length == number of samples
##' @param numeric flag determining class of return value, default FALSE
##' @param ... additional parameters passed to \code{snpStats::impute.snps}
##' @return matrix, either numeric (if numeric=TRUE) or SnpMatrix (if numeric=FALSE)
##' @author Chris Wallace
##' @export 
impute.missing <- function(X,bp=1:ncol(X),strata=NULL, numeric=FALSE, ...) {
  N <- as(X,"numeric")
  if(!is.null(strata)) {
    strata <- as.factor(strata)
    if(length(levels(strata))>10)
      stop("too many levels in strata\n")
    for(i in levels(strata)) {
      cat("\nstrata",i,"\n")
      wh <- which(strata==i)
      N[wh,] <- impute.missing(X[wh,,drop=FALSE],bp, numeric=TRUE, ...)
    }
  } else {
    csumm <- col.summary(X)
    use <- csumm[,"Certain.calls"]==1
    X2 <- X[,use]
    bp <- bp[use]
    imp <- (csumm[,"Call.rate"]<1)[use]
    cat(sum(imp),"to impute\n")
    for(i in which(imp)) {
      cat(i,".")
      rule <- snp.imputation(X2[,-i,drop=FALSE],X2[,i,drop=FALSE],bp[-i],bp[i])
      if(is.null(rule@.Data[[1]]))
        next
      imp <- impute.snps(rules=rule,snps=X2[,rule@.Data[[1]]$snps,drop=FALSE], ...)
      wh.na <- which(is.na(N[,i]))
      N[wh.na, colnames(X2)[i]] <- imp[wh.na]
    }
    cat("\n")
  }
  if(numeric) {
    return(N)
  } else {
    N <- pmin(N,2)
    N <- pmax(N,0)
    if(is(X,"aSnpMatrix")) {
      return(new("aSnpMatrix",
                 .Data=suppressWarnings(new("SnpMatrix",data=round(N)+1, nrow=nrow(N), ncol=ncol(N), dimnames=dimnames(N))),
                 snps=X@snps,
                 samples=X@samples,
                 phenotype=X@phenotype,
                 alleles=X@alleles))
    } else {
      return(new("SnpMatrix",
                 .Data=suppressWarnings(new("SnpMatrix",data=round(N)+1, nrow=nrow(N), ncol=ncol(N), dimnames=dimnames(N)))))
    }
  }
}