R/predict.geneJAM.R

In abuchardt/EdGwas: Joint regression Analysis of Multiple traits based on GENEtic relationships

#' Extract coefficients from an geneJAM object
#'
#' Similar to other predict methods, this function predicts fitted values and extract coefficients from a fitted "geneJAM" object.
#'
#' @param object Fitted "geneJAM" object.
#' @param rho Value(s) of the penalty parameter rho at which predictions are required. Default (NULL) is the entire sequence used to create the model.
#' @param newPS Matrix of new values for x of dimension nobs x nouts (must be polygenic scores - can be generated by \code{\link{ps.geneJAM}}) at which predictions are to be made. This argument is not used for \code{type = "coefficients"}.
#' @param type Type of prediction required. Type "link" returns a list of length nouts of matrices of dimension nobs x length(rho) of the fitted values for each response. Type "coefficients" gives the coefficients at the requested values for rho.
#' @param ... additional arguments affecting the predictions produced.
#'
#' @return The object returned depends on type.
#'
#' @examples
#' N <- 1000 #
#' q <- 10 #
#' p <- 1000 #
#' set.seed(1)
#' # Sample 1
#' x0 <- matrix(rbinom(n = N*p, size = 2, prob = 0.3), nrow=N, ncol=p)
#' B <- matrix(0, nrow = p, ncol = q)
#' B[1, 1:2] <- 2.5
#' y0 <- x0 %*% B + matrix(rnorm(N*q), nrow = N, ncol = q)
#' beta <- ps.geneJAM(x0, y0)$beta
#' # Sample 2
#' x <- matrix(rbinom(n = N*p, size = 2, prob = 0.3), nrow=N, ncol=p)
#' y <- x %*% B + matrix(rnorm(N*q), nrow = N, ncol = q)
#' ps <- x %*% beta
#' ###
#' pc <- geneJAM(ps[-(1:100), ], y[-(1:100), ])
#' newy <- predict(pc, newPS = ps[(1:100), ], rho = pc$rho.min)[[1]]
#'
#' mean((newy[,1] - y[1:100, 1])^2)
#' pss <- ps[-(1:100), 1]
#' lmfit <- lm(y[-(1:100), 1] ~ pss)
#' newylm <- predict(lmfit, newdata = data.frame(pss = ps[(1:100), 1]))
#' mean((newylm - y[1:100, 1])^2)
#'
#' @export
#'
predict.geneJAM <- function(object, newPS, rho = NULL,
                           type = c("link", "coefficients"), #"response","nonzero","class"),
                           #exact = FALSE, newoffset
                           ...
                           ){
  type <- match.arg(type)
  if (missing(newPS)){
    if (!match(type, "coefficients", FALSE))
      stop("Value for 'newPS' missing")
  }

  if (is.null(rho)){
    rho <- object$rho
  } else if (is.numeric(rho)) {
    if (!all(rho %in% object$rho))
      stop("rho needs to be a value used to create the model")
  } else stop("Invalid form for rho")

  nrho <- length(rho)
  whichRho <- match(rho, object$rho, FALSE)
  nouts <- ncol(newPS)

  # Simple lm predictions
  #nlmfit <- matrix(NA, nrow(newPS), nouts)
  #for (l in seq(nouts)) {
  #  xl <- newPS[, l]
  #  nlmfit[,l] <- predict(object$lmfit[[l]],
  #                  newdata = data.frame(xl))
  #}




    ncoef <- lapply(seq(nouts), FUN = function(l) {
      out <- rbind(object$xi[l, whichRho], object$beta[l, whichRho])
      colnames(out) <- paste0("rho", whichRho)
      rownames(out) <- c("xi", "beta")
      out
    })
    names(ncoef) <- paste0("y", seq(nouts))

    if (type == "coefficients")
      return(ncoef)

    nfit0 <- vector(mode = "list", length = nrho)
    nfit0 <- lapply(nfit0, FUN = function(l) matrix(NA, nrow(newPS), nouts))


  nfit <- vector(mode = "list", length = nrho)
  nfit <- lapply(nfit, FUN = function(l) matrix(NA, nrow(newPS), nouts))
  rhoseq <- seq(nrho)

  for (j in rhoseq) {


      for (l in seq(nouts)) {
        nfit[[j]][, l] <- cbind(1, newPS[, l]) %*% ncoef[[l]][, j]
      }

    }

  nfit
}