hack to get time, yhat, lineage from tradeSeq for one gene

getCurves.R

getCurves <- function(models, counts, gene, nPoints = 100, lwd = 2, size = 2/3, alpha = 2/3) {
  if (is(gene, "character")) {
    id <- which(names(models) %in% gene)
  }
  else id <- gene
  dm <- colData(models)$tradeSeq$dm
  y <- unname(counts[id, ])
  X <- colData(models)$tradeSeq$X
  slingshotColData <- colData(models)$slingshot
  pseudotime <- slingshotColData[, grep(x = colnames(slingshotColData),
                                        pattern = "pseudotime")]
  if (is.null(dim(pseudotime)))
    pseudotime <- matrix(pseudotime, ncol = 1)
  nCurves <- length(grep(x = colnames(dm), pattern = "t[1-9]"))
  betaMat <- rowData(models)$tradeSeq$beta[[1]]
  beta <- betaMat[id, ]
  lcol <- timeAll <- rep(0, nrow(dm))
  for (jj in seq_len(nCurves)) {
    for (ii in seq_len(nrow(dm))) {
      if (dm[ii, paste0("l", jj)] == 1) {
        timeAll[ii] <- dm[ii, paste0("t", jj)]
        lcol[ii] <- jj
      }
      else {
        next
      }
    }
  }
  df <- data.frame(time = timeAll, gene_count = y,
                   lineage = as.character(lcol))
  rows <- sample(seq_len(nrow(df)), nrow(df), replace = FALSE)
  df <- df[rows, ]
  out <- list()
  for (jj in seq_len(nCurves)) {
    df <- tradeSeq:::.getPredictRangeDf(dm, jj, nPoints = nPoints)
    Xdf <- tradeSeq:::predictGAM(lpmatrix = X, df = df, pseudotime = pseudotime)
    yhat <- c(exp(t(Xdf %*% t(beta)) + df$offset))
    out[[jj]] <- data.frame(time=df[,paste0("t",jj)], yhat, lineage=as.character(jj))
  }
  return(do.call(rbind, out))
}