R/interpolate.R

In vital: Tidy Analysis Tools for Mortality, Fertility, Migration and Population Data

#' Interpolate missing values using a vital model
#'
#' Uses a fitted vital model to interpolate missing values from a dataset.
#'
#' @param object A mable containing a single model column.
#' @param new_data A dataset with the same structure as the data used to fit the model.
#' @param ... Other arguments passed to interpolate methods.
#' @return A vital object with missing values interpolated.
#' @examples
#' nor_female <- norway_mortality |>
#'   dplyr::filter(Sex == "Female")
#' nor_female |>
#'   model(mean = FMEAN(Mortality)) |>
#'   interpolate(nor_female)
#' @rdname interpolate
#' @author Rob J Hyndman
#' @export
interpolate.mdl_vtl_df <- function(object, new_data, ...) {
  if (length(mable_vars(object)) > 1) {
    abort(
      "Interpolation can only be done using one model. \nPlease use select() to choose the model to interpolate with."
    )
  }
  keys <- key_vars(new_data)
  agevar <- age_var(new_data)
  keys_noage <- keys[keys != agevar]
  index <- index_var(new_data)
  object <- bind_new_data(object, new_data)
  object <- transmute(
    as_tibble(object),
    !!!syms(keys_noage),
    interpolated = map2(!!sym(mable_vars(object)), new_data, interpolate, ...)
  )
  unnest_tbl(object, "interpolated") |>
    as_tsibble(index = index, key = all_of(c(agevar, keys_noage))) |>
    as_vital(.age = agevar, reorder = TRUE)
}

#' @export
interpolate.mdl_vtl_ts <- function(object, new_data, ...) {
  object$model$stage <- "interpolate"
  object$model$add_data(new_data)
  specials <- tryCatch(
    parse_model_rhs(object$model),
    error = function(e) {
      abort(sprintf(
        "%s\nUnable to compute required variables from provided `new_data`.\nDoes your interpolation data include all variables required by the model?",
        e$message
      ))
    },
    interrupt = function(e) {
      stop("Terminated by user", call. = FALSE)
    }
  )
  object$model$remove_data()
  object$model$stage <- NULL
  resp <- map2(
    seq_along(object$response),
    object$response,
    function(i, resp) {
      expr(object$transformation[[!!i]](!!resp))
    }
  ) %>%
    set_names(map_chr(object$response, as_string))
  vvar <- vital_vars(new_data)
  agevar <- vvar["age"]
  age <- new_data[[agevar]]
  new_data <- transmute(new_data, !!!resp)
  new_data[[agevar]] <- age
  attr(new_data, "vital") <- vvar
  new_data <- interpolate(
    object[["fit"]],
    new_data = new_data,
    specials = specials,
    ...
  )
  new_data[names(resp)] <- map2(
    new_data[names(resp)],
    object$transformation,
    function(x, f) invert_transformation(f)(x)
  )
  new_data
}