R/too_far.R

In gratia: Graceful 'ggplot'-Based Graphics and Other Functions for GAMs Fitted Using 'mgcv'

#' Set rows of data to `NA` if the lie too far from a reference set of values
#'
#' @param smooth an mgcv smooth object
#' @param input data frame containing the input observations and the columns to
#'   be set to `NA`
#' @param reference data frame containing the reference values
#' @param cols character vector of columns whose elements will be set to `NA` if
#'   the data lies too far from the reference set
#' @param dist numeric, the distance from the reference set beyond which
#'   elements of `input` will be set to `NA`
#'
#' @export
`too_far_to_na` <- function(smooth, input, reference, cols, dist = NULL) {
    UseMethod("too_far_to_na")
}

#' @export
#' @importFrom dplyr mutate across
#' @importFrom tidyselect all_of
`too_far_to_na.mgcv.smooth` <- function(smooth, input, reference, cols,
                                        dist = NULL) {
    # only for smooths of 2D currently
    sm_dim <- smooth_dim(smooth)
    if (sm_dim < 2L || sm_dim > 3L) {
        return(input)
    }

    # call too_far to identify the observations too far from reference
    input <- .call_too_far(smooth = smooth, input = input,
                           reference = reference, cols = cols, dist = dist)

    # return
    input
}

#' @export
#' @importFrom dplyr mutate across
#' @importFrom tidyselect all_of
`too_far_to_na.t2.smooth` <- function(smooth, input, reference, cols,
                                      dist = NULL) {
    # only for smooths of 2D currently
    sm_dim <- smooth_dim(smooth)
    if (sm_dim < 2L || sm_dim > 3L) {
        return(input)
    }

    # call too_far to identify the observations too far from reference
    input <- .call_too_far(smooth = smooth, input = input,
                           reference = reference, cols = cols, dist = dist)

    # return
    input
}

#' @export
#' @importFrom dplyr mutate across
#' @importFrom tidyselect all_of
`too_far_to_na.tensor.smooth` <- function(smooth, input, reference, cols,
                                          dist = NULL) {
    # only for smooths of 2D currently
    sm_dim <- smooth_dim(smooth)
    if (sm_dim < 2L || sm_dim > 3L) {
        return(input)
    }

    # call too_far to identify the observations too far from reference
    input <- .call_too_far(smooth = smooth, input = input,
                           reference = reference, cols = cols, dist = dist)

    # return
    input
}

#' Exclude values that lie too far from the support of data
#'
#' Identifies pairs of covariate values that lie too far from the original data.
#' The function is currently a basic wrapper around [mgcv::exclude.too.far()].
#'
#' @param x,y numeric; vector of values of the covariates to compare with
#'   the observed data
#' @param ref_1,ref_2 numeric; vectors of covariate values that represent the
#'   reference against which `x1 and `x2` are compared
#' @param dist if supplied, a numeric vector of length 1 representing the
#'   distance from the data beyond which an observation is excluded. For
#'   example, you want to exclude values that lie further from an observation
#'   than 10% of the range of the observed data, use `0.1`.
#'
#' @return Returns a logical vector of the same length as `x1`.
#'
#' @importFrom mgcv exclude.too.far
#' @export
`too_far` <- function(x, y, ref_1, ref_2, dist = NULL) {
    ind <- if (is.null(dist)) {
        # If `NULL` keep everything == vector of TRUEs of correct length
        rep(TRUE, length.out = length(x))
    } else {
        # if `dist` is provided, check it is of the correct kind
        if (!is.numeric(dist) || !identical(length(dist), 1L)) {
            stop("'dist', if provided, must be a single numeric value.",
                 call. = FALSE)
        }
        # call exclude.too.far
        mgcv::exclude.too.far(x, y, ref_1, ref_2, dist = dist)
    }

    # return
    ind
}


#' Sets the elements of vector to `NA`
#'
#' Given a vector `i` indexing the elements of `x`, sets the selected elements
#' of `x` to `NA`.
#' 
#' @param x vector of values
#' @param i vector of values used to subset `x`
#' 
#' @return Returns `x` with possibly some elements set to `NA`
#' 
#' @export
`to_na` <- function(x, i) {
    # check x and i are of the same length
    if (identical(length(x), length(i))) {
        # set the indicated elements of `x` to `NA`
        x[i] <- NA
    } else {
        stop("'x' and 'i' must be the same legnth.")
    }
    x
}

#' Set up and call `too_far` on the supplied input
#'
#' @param smooth an mgcv smooth object
#' @param input data frame containing the input observations and the columns to
#'   be set to `NA`
#' @param reference data frame containing the reference values
#' @param cols character vector of columns whose elements will be set to `NA` if
#'   the data lies too far from the reference set
#' @param dist numeric, the distance from the reference set beyond which
#'   elements of `input` will be set to `NA`
#' 
#' @noRd
#' 
#' @keywords internal
`.call_too_far` <- function(smooth, input, reference, cols, dist) {
    # Handle the case where `input` is a nested tibble & data are in $data
    was_nested <- FALSE
    if (!is.null(input[["data"]]) &&
        (inherits(input, "tbl_df") && is.list(input[["data"]]))) {
        input <- unnest(input, cols = all_of("data"))
        was_nested <- TRUE
    }
    # what variables do we need to work on
    sm_vars <- smooth_variable(smooth)
    # indicator for observations too far from grid
    ind <- too_far(x = input[[sm_vars[1L]]],
                   y = input[[sm_vars[2L]]],
                   ref_1 = reference[[sm_vars[1L]]],
                   ref_2 = reference[[sm_vars[2L]]],
                   dist = dist)
    # set the indicated columns to `NA`
    input <- mutate(input,
        across(all_of(cols), \(x) to_na(x, i = ind)))

    # nest `input` again if it was nested previously
    if (was_nested) {
        input <- nest(input, data = !all_of(c("smooth", "type", "by")))
    }

    # return
    input
}