R/methods_flexsurv.R

In marginaleffects: Predictions, Comparisons, Slopes, Marginal Means, and Hypothesis Tests

#' @rdname set_coef
#' @export
set_coef.flexsurvreg <- function(model, coefs, ...) {
  out <- model
  out$res[, 1] <- coefs
  out$coefficients <- coefs
  return(out)
}


#' @rdname get_predict
#' @export
get_predict.flexsurvreg <- function(model, newdata, type, ...) {
  preds <- stats::predict(
    object = model,
    newdata = newdata,
    type = type,
    ...
  )

  if (ncol(preds) == 1L) {
    if (names(preds) == '.pred') {
      gp <- unlist(lapply(preds$.pred, function(x) {x[, 1, drop = TRUE]}))
      val <- unlist(lapply(preds$.pred, function(x) {x[, 2, drop = TRUE]}))

      out <- data.frame(
        rowid = seq_len(nrow(preds)),
        group = as.vector(gp),
        estimate = as.vector(val)
      )
      out$group <- group_to_factor(out$group, model)
      return(out)
    }
    out <- data.frame(
      rowid = seq_len(nrow(preds)),
      estimate = as.vector(preds[, 1, drop = TRUE])
    )
    return(out)
  }

  out <- data.frame(
    rowid = seq_len(nrow(preds)),
    group = as.vector(preds[, 1, drop = TRUE]),
    estimate = as.vector(preds[, 2, drop = TRUE])
  )
  out$group <- group_to_factor(out$group, model)
  out
}