R/make-newdata.R

In pammtools: Piece-Wise Exponential Additive Mixed Modeling Tools for Survival Analysis

#' Extract information of the sample contained in a data set
#'
#' Given a data set and grouping variables, this function returns mean values
#' for numeric variables and modus for characters and factors. Usually
#' this function should not be called directly but will rather be called
#' as part of a call to \code{make_newdata}.
#'
#' @rdname sample_info
#' @param x A data frame (or object that inherits from \code{data.frame}).
#' @importFrom stats median
#' @return A data frame containing sample information (for each group).
#' If applied to an object of class \code{ped}, the sample means of the
#' original data is returned.
#' Note: When applied to a \code{ped} object, that doesn't contain covariates
#' (only interval information), returns data frame with 0 columns.
#'
#' @export
#' @keywords internal
sample_info <- function(x) {
  UseMethod("sample_info", x)
}

#' @import checkmate dplyr
#' @importFrom purrr compose
#' @export
#' @rdname sample_info
sample_info.data.frame <- function(x) {
  cn <- colnames(x)
  num <- summarize_if(x, .predicate = is.numeric, ~ mean(., na.rm = TRUE))
  fac <- summarize_if(x, .predicate = compose("!", is.numeric), modus)

  nnames <- intersect(names(num), names(fac))

  if (length(nnames) != 0) {
    suppressMessages(
      x <- left_join(num, fac) %>% group_by(!!!lapply(nnames, as.name))
    )
  } else {
    x <- bind_cols(num, fac)
  }

  return(select(x, one_of(cn)))
}

#' @rdname sample_info
#' @import checkmate dplyr
#' @importFrom rlang sym
#' @export
sample_info.ped <- function(x) {
  # is.grouped_df
  # remove "noise" information on interval variables
  grps <- group_vars(x)
  iv <- attr(x, "intvars")
  id_var <- attr(x, "id_var")
  if (!is.null(id_var) && id_var %in% names(x)) {
    x <- x %>%
      group_by(!!sym(id_var)) %>%
      slice(1) %>%
      ungroup()
  }
  x <- x %>%
    grouped_df(grps) %>%
    select(-one_of(iv))
  if (test_data_frame(x, min.rows = 1, min.cols = 1)) {
    sample_info.data.frame(x)
  } else {
    NULL
  }
}

#' @rdname sample_info
#' @export
sample_info.fped <- function(x) {
  x %>% select_if(~ !is.matrix(.x)) %>% sample_info.ped()
}


#' Create a data frame from all combinations of data frames
#'
#' Works like \code{\link[base]{expand.grid}} but for data frames.
#'
#' @importFrom dplyr slice bind_cols
#' @importFrom vctrs vec_c
#' @importFrom purrr map map_lgl
#' @importFrom checkmate test_data_frame
#' @param ... Data frames that should be combined to one data frame.
#' Elements of first df vary fastest, elements of last df vary slowest.
#' @examples
#' combine_df(
#'   data.frame(x=1:3, y=3:1),
#'   data.frame(x1=c("a", "b"), x2=c("c", "d")),
#'   data.frame(z=c(0, 1)))
#' @export
#' @keywords internal
combine_df <- function(...) {
  dots <- list(...)
  if (!all(sapply(dots, test_data_frame))) {
    stop("All elements in ... must inherit from data.frame!")
  }
  ind_seq <- map(dots, ~ seq_len(nrow(.x)))
  not_empty <- map_lgl(ind_seq, ~ length(.x) > 0)

  ord <- lapply(dots[not_empty], function(z) colnames(z)) |> unlist()
  out <- do.call(expand_grid, rev(dots[not_empty]))
  out <- out[, ord]
}


#' Construct a data frame suitable for prediction
#'
#' This functions provides a flexible interface to create a data set that
#' can be plugged in as \code{newdata} argument to a suitable  \code{predict}
#' function (or similar).
#' The function is particularly useful in combination with one of the
#' \code{add_*} functions, e.g., \code{\link[pammtools]{add_term}},
#' \code{\link[pammtools]{add_hazard}}, etc.
#'
#' @rdname newdata
#' @aliases make_newdata
#' @importFrom tidyr expand_grid
#' @inheritParams sample_info
#' @param ... Covariate specifications (expressions) that will be evaluated
#' by looking for variables in \code{x}. Must be of the form \code{z = f(z)}
#' where \code{z} is a variable in the data set and \code{f} a known
#' function that can be usefully applied to \code{z}. Note that this is also
#' necessary for single value specifications (e.g. \code{age = c(50)}).
#' For data in PED (piece-wise exponential data) format, one can also specify
#' the time argument, but see "Details" an "Examples" below.
#' @import dplyr
#' @importFrom checkmate assert_data_frame assert_character
#' @importFrom purrr map
#' @importFrom tidyr expand_grid
#' @details Depending on the type of variables in \code{x}, mean or modus values
#' will be used for variables not specified in ellipsis
#' (see also \code{\link[pammtools]{sample_info}}). If \code{x} is an object
#' that inherits from class \code{ped}, useful data set completion will be
#' attempted depending on variables specified in ellipsis. This is especially
#' useful, when creating a data set with different time points, e.g. to
#' calculate survival probabilities over time (\code{\link[pammtools]{add_surv_prob}})
#' or to calculate a time-varying covariate effects (\code{\link[pammtools]{add_term}}).
#' To do so, the time variable has to be specified in \code{...}, e.g.,
#' \code{tend = seq_range(tend, 20)}. The problem with this specification is that
#' not all values produced by \code{seq_range(tend, 20)} will be actual values
#' of \code{tend} used at the stage of estimation (and in general, it will
#' often be tedious to specify exact \code{tend} values). \code{make_newdata}
#' therefore finds the correct interval and sets \code{tend} to the respective
#' interval endpoint. For example, if the intervals of the PED object are
#' \eqn{(0,1], (1,2]} then \code{tend = 1.5} will be set to \code{2}.
#'
#' The returned data frame contains \code{tend}, \code{id}, the user-supplied
#' covariates (and \code{cause}/\code{transition} for competing risks /
#' multi-state models). Internal PED columns \code{tstart}, \code{intlen},
#' \code{interval}, \code{offset}, and \code{ped_status} are dropped.
#' Downstream \code{add_*} functions reconstruct \code{intlen} on demand via
#' \code{reconstruct_intlen()} when needed.
#' See examples below.
#' @examples
#' # General functionality
#' tumor %>% make_newdata()
#' tumor %>% make_newdata(age=c(50))
#' tumor %>% make_newdata(days=seq_range(days, 3), age=c(50, 55))
#' tumor %>% make_newdata(days=seq_range(days, 3), status=unique(status), age=c(50, 55))
#' # mean/modus values of unspecified variables are calculated over whole data
#' tumor %>% make_newdata(sex=unique(sex))
#' tumor %>% group_by(sex) %>% make_newdata()
#'
#' # Examples for PED data
#' ped <- tumor %>% slice(1:3) %>% as_ped(Surv(days, status)~., cut = c(0, 500, 1000))
#' ped %>% make_newdata(age=c(50, 55))
#'
#' # if time information is specified, other time variables will be specified
#' # accordingly and offset calculated correctly
#' ped %>% make_newdata(tend = c(1000), age = c(50, 55))
#' ped %>% make_newdata(tend = unique(tend))
#' ped %>% group_by(sex) %>% make_newdata(tend = unique(tend))
#'
#' # tend is set to the end point of respective interval:
#' ped <- tumor %>% as_ped(Surv(days, status)~.)
#' seq_range(ped$tend, 3)
#' make_newdata(ped, tend = seq_range(tend, 3))
#' @export
make_newdata <- function(x, ...) {
  UseMethod("make_newdata", x)
}


#' @inherit make_newdata
#' @rdname newdata
#' @export
make_newdata.default <- function(x, ...) {
  assert_data_frame(x, all.missing = FALSE, min.rows = 2, min.cols = 1)

  orig_names <- names(x)

  expressions <- quos(...)
  expr_evaluated <- map(expressions, lazyeval::f_eval, data = x) |>
    map(c)

  if ("tend" %in% names(expressions)) {
    if (any(expr_evaluated[["tend"]] <= 0)) {
      stop("'tend' variable must take values > 0.")
    }
  }

  # construct data parts depending on input type
  lgl_atomic <- map_lgl(expr_evaluated, is_atomic)
  # part1 <- expr_evaluated[lgl_atomic] |> cross_df()
  part1 <- do.call(tidyr::expand_grid, rev(expr_evaluated[lgl_atomic]))
  part2 <- do.call(combine_df, expr_evaluated[!lgl_atomic])

  ndf <- combine_df(part1, part2)
  rest <- x %>% select(-one_of(c(colnames(ndf))))
  if (ncol(rest) > 0) {
    si <- sample_info(rest) %>% ungroup()
    ndf <- combine_df(si, ndf)
  }

  ndf %>% select(one_of(orig_names))
}

#' @rdname newdata
#' @inherit make_newdata.default
#' @export
make_newdata.ped <- function(x, ...) {
  assert_data_frame(x, all.missing = FALSE, min.rows = 2, min.cols = 1)

  # prediction time points have to be interval end points so that piece-wise
  # constancy of predicted hazards is respected. If user overrides this, warn.
  orig_vars <- names(x)
  int_df <- int_info(x)
  brks <- int_df$tend # get original estimation cut points

  expressions <- quos(...)
  dot_names <- names(expressions)
  int_names <- names(int_df)

  # x <- select(x, -one_of(setdiff(int_names, c(dot_names, "intlen", "intmid"))))
  ndf <- x %>%
    select(-one_of(setdiff(int_names, c(dot_names, "intlen", "intmid")))) %>%
    unped() %>%
    make_newdata(...)

  if (any(names(int_df) %in% names(ndf))) {
    times <- ndf$tend
    ped_times <- sort(unique(union(c(0, brks), times)))
    # extract relevant intervals only, keeps data small
    ped_times <- ped_times[ped_times <= max(times)]

    info <- get_intervals(x, times)
    int_tend <- info$tend
    if (!all(ndf$tend == int_tend)) {
      message("Not all requested timepoints correspond to original cut points.")
    }
    ndf$tend <- int_tend
    map_times <- data.frame(
      tend = sort(unique(info$times)),
      tstart_lag = lag(sort(unique(info$times)), default = 0)
    )
    suppressMessages(
      ndf <- ndf %>%
        left_join(int_df) %>%
        mutate(tend = info$times) %>%
        left_join(map_times) %>%
        mutate(tstart = pmax(.data$tstart, .data$tstart_lag)) %>%
        select(-one_of("tstart_lag")) # correct tend and tstart
    )
  } else {
    ndf <- combine_df(int_df[1, ], ndf)
  }

  int_names <- intersect(int_names, c("intlen", orig_vars))
  ndf <- ndf %>%
    select(one_of(c(int_names, setdiff(orig_vars, int_names))))

  # drop internal columns not needed for prediction
  drop_vars <- c("tstart", "intlen", "interval", "offset", "ped_status")
  ndf <- ndf %>% select(-any_of(drop_vars))

  # needed for expand df
  attr(ndf, "trafo_args") <- attr(x, "trafo_args")
  attr(ndf, "intvars") <- attr(x, "intvars")

  ndf
}

#' @rdname newdata
#' @inherit make_newdata.ped
#' @importFrom rlang quos
#' @export
make_newdata.fped <- function(x, ...) {
  assert_data_frame(x, all.missing = FALSE, min.rows = 2, min.cols = 1)

  # prediction time points have to be interval end points so that piece-wise
  # constancy of predicted hazards is respected. If user overrides this, warn.
  expressions <- quos(...)
  dot_names <- names(expressions)
  cumu_vars <- setdiff(unlist(attr(x, "func_mat_names")), dot_names)
  cumu_smry <- smry_cumu_vars(x, attr(x, "time_var")) %>%
    select(one_of(cumu_vars))

  ndf <- x %>%
    select(one_of(setdiff(names(x), cumu_vars))) %>%
    unfped() %>%
    make_newdata(...)

  out_df <- do.call(combine_df, compact(list(cumu_smry, ndf)))
  int_df <- int_info(attr(x, "breaks"))
  suppressMessages(
    out_df <- out_df %>%
      left_join(int_df) %>%
      select(-one_of(c("intmid"))) %>%
      as_tibble()
  )

  # adjust lag-lead indicator
  out_df <- adjust_ll(out_df, x)

  # drop internal columns not needed for prediction
  drop_vars <- c("tstart", "intlen", "interval", "offset", "ped_status")
  out_df <- out_df %>% select(-any_of(drop_vars))

  out_df
}


smry_cumu_vars <- function(data, time_var) {
  cumu_vars <- unlist(attr(data, "func_mat_names"))
  func_list <- attr(data, "func")
  z_vars <- map(func_list, ~ get_zvars(.x, time_var, length(func_list))) %>%
    unlist()
  smry_z <- select(data, one_of(z_vars)) %>%
    map(~ .x[1, ]) %>%
    map(~ mean(unlist(.x))) %>%
    bind_cols()
  smry_time <- select(data, setdiff(cumu_vars, z_vars)) %>% map(~ .x[1, 1])

  bind_cols(smry_z, smry_time)
}

get_zvars <- function(func, time_var, n_func) {
  col_vars <- func$col_vars
  all_vars <- make_mat_names(
    c(col_vars, "LL"),
    func$latency_var,
    func$tz_var,
    func$suffix,
    n_func
  )
  time_vars <- make_mat_names(
    c(time_var, func$tz_var, "LL"),
    func$latency_var,
    func$tz_var,
    func$suffix,
    n_func
  )

  setdiff(all_vars, time_vars)
}


## apply ll_fun to newly created data
adjust_ll <- function(out_df, data) {
  func_list <- attr(data, "func")
  n_func <- length(func_list)
  LL_names <- grep("LL", unlist(attr(data, "func_mat_names")), value = TRUE)
  time_var <- attr(data, "time_var")
  if (is.null(time_var) || !time_var %in% colnames(out_df)) {
    time_var <- "tend"
  }

  for (i in LL_names) {
    ind_ll <- map_lgl(names(attr(data, "ll_funs")), ~ grepl(.x, i))
    if (any(ind_ll)) {
      ind_ll <- which(ind_ll)
    } else {
      ind_ll <- 1
    }

    func <- func_list[[ind_ll]]
    ll_i <- attr(data, "ll_funs")[[ind_ll]]
    tz_var <- attr(data, "tz_vars")[[ind_ll]]
    tz_var <- make_mat_names(
      tz_var,
      func$latency_var,
      func$tz_var,
      func$suffix,
      n_func
    )
    if (func$latency_var == "") {
      out_df[[i]] <- ll_i(out_df[[time_var]], out_df[[tz_var]]) * 1L
    } else {
      out_df[[i]] <- ll_i(
        out_df[[time_var]],
        out_df[[time_var]] -
          out_df[[tz_var]]
      ) *
        1L
    }
  }

  out_df
}


#' Reconstruct intlen from time variable and stored cut points
#'
#' Computes interval lengths from the sorted unique values of the time variable
#' in newdata. This is used by add_* functions that need intlen for cumulative
#' calculations. If \code{tstart} is not available, the first interval length is
#' taken as the first sorted time value (implicitly assuming a 0-origin time
#' scale).
#' @param newdata A data frame with a time column (default \code{tend}).
#' @param time_var Character name of the time variable. Defaults to
#'   \code{"tend"}.
#' @param interval_length Character name of the interval-length column to create.
#' @return The input data frame with an \code{intlen} column added.
#' @keywords internal
reconstruct_intlen <- function(
  newdata,
  time_var = "tend",
  interval_length = "intlen"
) {
  if (interval_length %in% colnames(newdata)) return(newdata)

  assert_data_frame(newdata, all.missing = FALSE)
  assert_string(time_var)
  assert_choice(time_var, colnames(newdata))
  assert_string(interval_length)
  assert_numeric(newdata[[time_var]], any.missing = FALSE)

  if ("tstart" %in% colnames(newdata)) {
    assert_numeric(newdata[["tstart"]], any.missing = FALSE)
    newdata[[interval_length]] <- newdata[[time_var]] - newdata[["tstart"]]
  } else {
    times <- sort(unique(newdata[[time_var]]))
    intlen_map <- c(times[1], diff(times))
    time_ind <- match(newdata[[time_var]], times)
    newdata[[interval_length]] <- unname(intlen_map[time_ind])
  }

  if (anyNA(newdata[[interval_length]])) {
    stop(
      "Failed to reconstruct '",
      interval_length,
      "' from time variable '",
      time_var,
      "'."
    )
  }
  if (any(newdata[[interval_length]] < 0)) {
    stop(
      "Reconstructed '",
      interval_length,
      "' contains negative values for time variable '",
      time_var,
      "'."
    )
  }
  newdata
}


## apply expand df to complete make newdata for arbitrary time points
expand_df <- function(
  x,
  object,
  trafo_args = NULL,
  intvars = NULL,
  time_var = NULL,
  interval_length = "intlen",
  ...
) {

  orig_vars <- names(x)
  if (is.null(time_var)) {
    time_var <- attr(x, "time_var")
  }
  if (is.null(time_var)) {
    time_var <- "tend"
  }

  id_var <- trafo_args[["id"]]
  brks <- trafo_args[["cut"]]

  int_names <- setdiff(intvars, id_var)
  interval_struct_vars <- c(time_var, "tstart", interval_length, "offset", "ped_status")
  cov_names <- setdiff(names(x), c(id_var, interval_struct_vars))

  expressions <- quos(...)
  dot_names <- names(expressions)

  times <- unique(x[[time_var]])

  ped_times <- sort(unique(union(c(0, brks), times)))
  # extract relevant intervals only, keeps data small
  ped_times <- ped_times[ped_times <= max(times)]
  # obtain interval information
  ped_info <- get_intervals(brks, ped_times[-1])

  ped_info[[interval_length]] <- c(
    ped_info[["times"]][1],
    diff(ped_info[["times"]])
  )
  ped_info[[time_var]] <- ped_info[["times"]]

  if (length(cov_names) == 0) {
    newdata <- ped_info
  } else {
    newdata <- x %>%
      select(any_of(c(cov_names))) %>%
      distinct() # otherwise combine_df duplicates rows

    newdata <- combine_df(ped_info, newdata) %>% select(-c("times"))
  }

  map_times <- data.frame(
    time_tmp = sort(ped_times),
    tstart_lag = lag(ped_times, default = 0)
  )
  names(map_times)[1] <- time_var
  suppressMessages(
    newdata <- newdata %>%
      left_join(map_times, by = time_var) %>%
      mutate(
        tstart = pmax(.data$tstart, .data$tstart_lag)
      ) %>%
      select(-one_of("tstart_lag")) # correct tstart
  )
  newdata[[interval_length]] <- newdata[[time_var]] - newdata[["tstart"]]

  # if(length(haz_vars_in_data) != 0) {
  #   if(length(setdiff(haz_vars_in_data, vars_exclude)) != 0){
  #     newdata <- get_hazard(object, newdata, type = "response", ci = ci,
  #                ci_type = ci_type, time_var = time_var, se_mult = se_mult, ...)
  #   } else {
  #     newdata <- get_hazard(object, newdata, type = "response", ci = FALSE,
  #                time_var = time_var, ...)
  #   }
  # }

  newdata %>% select(any_of(c(orig_vars, interval_length)))
}

deduce_df <- function(x, object, ...) {
}

#' Ensure all breakpoints are present in newdata for cumulative calculations
#'
#' Checks whether all cut points up to the maximum observed time are present in
#' \code{newdata}. If not, expands the data frame to include the missing
#' breakpoints via \code{expand_df}. In either case the function guarantees that
#' an interval-length column (default \code{intlen}) exists on return.
#' Existing grouping is preserved after expansion.
#'
#' @param newdata A data frame with a time column and, optionally, grouping.
#'   Must carry the \code{trafo_args} and \code{intvars} attributes set by
#'   \code{as_ped}.
#' @param object A fitted PAM/PAMM model object, passed to \code{expand_df}.
#' @param time_var Character name of the time variable (e.g. \code{"tend"}).
#' @param interval_length Character name of the interval-length column. If
#'   absent from \code{newdata} it is created via \code{reconstruct_intlen}.
#' @return A data frame with all required breakpoints present and an
#'   \code{interval_length} column guaranteed to exist.
#' @keywords internal
reconstruct_cutpoints <- function(
    newdata,
    object,
    time_var,
    interval_length
) {
  trafo_args <- attr(newdata, "trafo_args")
  intvars    <- attr(newdata, "intvars")
  times      <- setdiff(sort(unique(newdata[[time_var]])), 0)
  brks       <- setdiff(trafo_args[["cut"]][trafo_args[["cut"]] <= max(times)], 0)
  
  if (all(brks %in% times)) {
    # all breakpoints already present, just ensure intlen exists
    if (!interval_length %in% colnames(newdata)) {
      newdata <- reconstruct_intlen(newdata, time_var = time_var,
                                    interval_length = interval_length)
    }
    return(newdata)
  }

  # need to expand -- respect existing groups
  old_groups <- dplyr::groups(newdata)
  expanded <- group_split(newdata) |>
    map(\(.x) expand_df(
      .x,
      object        = object,
      trafo_args    = trafo_args,
      intvars       = intvars,
      time_var      = time_var,
      interval_length = interval_length
    )) |>
    map(\(.x) if (length(old_groups)) group_by(.x, !!!old_groups) else .x) |>
    bind_rows()
  
  if (!interval_length %in% colnames(expanded)) {
    expanded <- reconstruct_intlen(expanded, time_var = time_var,
                                   interval_length = interval_length)
  }

  expanded
}


# All variables that represent follow-up time should have the same values
# adjust_time_vars <- function(out_df, data, dot_names) {

#   time_vars <- c("tend",
#     grep(attr(data, "time_var"), unlist(attr(data, "func_mat_names")), value=TRUE))
#   time_vars_dots <- c(grep("tend", dot_names, value=TRUE),
#     grep(attr(data, "time_var"), dot_names, value=TRUE))
#   if (length(time_vars_dots) == 0) {
#     time_vars_dots <- "tend"
#   } else {
#     if (length(time_vars_dots) > 1) {
#       warning(paste0("Only one of ", paste0(time_vars_dots, collapse=", "),
#         "must be specified. Only the first one will be used!"))
#       time_vars_dots <- time_vars_dots[1]
#     }
#   }
#   for (i in setdiff(time_vars, time_vars_dots)) {
#     out_df[[i]] <- out_df[[time_vars_dots]]
#   }

#   out_df

# }