R/fit_local_rf.R

In gwrf: Geographically Weighted Random Forests

#' Fit one local random forest
#'
#' @param data Data frame containing response and predictors.
#' @param formula Model formula.
#' @param coords Numeric matrix/data frame with 2 columns.
#' @param focal_index Integer index of the focal observation.
#' @param bandwidth Adaptive k or fixed distance threshold.
#' @param adaptive Logical; adaptive or fixed neighborhood.
#' @param kernel Kernel type.
#' @param num.trees Number of trees for ranger.
#' @param mtry Number of variables tried at each split.
#' @param min.node.size Minimum node size for ranger.
#' @param importance Importance type passed to ranger.
#' @param use_case_weights Logical; whether to pass kernel weights to ranger.
#' @param seed Optional random seed.
#' @param keep_model Logical; whether to retain fitted ranger model.
#' @param neighbor_unit Character string indicating whether neighborhoods are
#'   defined using individual data rows (\code{"row"}) or unique spatial
#'   locations (\code{"location"}).
#' @param location_id Optional vector identifying the spatial location
#'   associated with each observation. Required when
#'   \code{neighbor_unit = "location"}.
#'
#' @return A named list containing the results for one focal local random
#' forest:
#' \describe{
#'   \item{focal_index}{Integer index of the focal observation.}
#'   \item{observed}{Observed response value at the focal observation.}
#'   \item{prediction}{Prediction from the local random forest for the focal
#'   observation.}
#'   \item{residual}{Observed minus predicted response for the focal
#'   observation.}
#'   \item{n_local}{Number of complete observations used to fit the local
#'   random forest.}
#'   \item{local_bandwidth}{Realized spatial bandwidth of the local
#'   neighborhood.}
#'   \item{variable_importance}{Named numeric vector of local variable
#'   importance values, or `NA` when importance cannot be calculated.}
#'   \item{ranger_model}{The fitted `ranger` model when `keep_model = TRUE`;
#'   otherwise `NULL`.}
#' }
#'
#' @examples
#' set.seed(1)
#'
#' n <- 20
#' dat <- data.frame(
#'   y = rnorm(n),
#'   x1 = rnorm(n),
#'   x2 = runif(n)
#' )
#' coords <- cbind(seq_len(n), rep(0, n))
#'
#' local_fit <- fit_local_rf(
#'   data = dat,
#'   formula = y ~ x1 + x2,
#'   coords = coords,
#'   focal_index = 10,
#'   bandwidth = 12,
#'   adaptive = TRUE,
#'   num.trees = 10,
#'   seed = 1
#' )
#'
#' local_fit$prediction
#' @export
fit_local_rf <- function(
    data,
    formula,
    coords,
    focal_index,
    bandwidth,
    adaptive = TRUE,
    kernel = "bisquare",
    num.trees = 500,
    mtry = NULL,
    min.node.size = 5,
    importance = "permutation",
    use_case_weights = TRUE,
    seed = NULL,
    keep_model = FALSE,
    neighbor_unit = "row",
    location_id = NULL
) {
  if (!is.null(seed)) {
    set.seed(seed)
  }

  nb <- build_neighbors(
    coords = coords,
    focal_index = focal_index,
    bandwidth = bandwidth,
    adaptive = adaptive,
    neighbor_unit = neighbor_unit,
    location_id = location_id
  )

  idx <- nb$neighbor_index
  local_data <- data[idx, , drop = FALSE]

  local_weights <- kernel_weights(
    distances = nb$distances,
    bandwidth = max(nb$local_bandwidth, .Machine$double.eps),
    kernel = kernel
  )

  # Ensure focal point has nonzero weight
  if (focal_index %in% idx) {
    focal_local_pos <- match(focal_index, idx)
    local_weights[focal_local_pos] <- max(local_weights[focal_local_pos], 1)
  }

  # Remove rows with missing model terms
  model_vars <- all.vars(formula)

  keep_rows <- complete.cases(local_data[, model_vars, drop = FALSE])

  local_data_complete <- local_data[keep_rows, , drop = FALSE]
  local_weights_complete <- local_weights[keep_rows]


  if (nrow(local_data_complete) < 10) {
    return(list(
      focal_index = focal_index,
      observed = NA_real_,
      prediction = NA_real_,
      residual = NA_real_,
      n_local = nrow(local_data_complete),
      local_bandwidth = nb$local_bandwidth,
      variable_importance = NA,
      ranger_model = NULL
    ))
  }

  if (is.null(mtry)) {
    predictor_names <- attr(stats::terms(formula), "term.labels")
    mtry <- max(1, floor(sqrt(length(predictor_names))))
  }

  ranger_args <- list(
    formula = formula,
    data = local_data_complete,
    num.trees = num.trees,
    mtry = mtry,
    min.node.size = min.node.size,
    importance = importance,
    write.forest = TRUE,
    keep.inbag = FALSE,
    respect.unordered.factors = "order"
  )

  if (use_case_weights) {
    ranger_args$case.weights <- local_weights_complete
  }

  fit <- do.call(ranger::ranger, ranger_args)

  focal_row <- data[focal_index, , drop = FALSE]
  pred <- predict(fit, data = focal_row)$predictions

  observed <- model.response(model.frame(formula, data = focal_row))
  residual <- observed - pred

  vi <- fit$variable.importance
  if (is.null(vi)) {
    vi <- NA
  }

  list(
    focal_index = focal_index,
    observed = as.numeric(observed),
    prediction = as.numeric(pred),
    residual = as.numeric(residual),
    n_local = nrow(local_data_complete),
    local_bandwidth = nb$local_bandwidth,
    variable_importance = vi,
    ranger_model = if (keep_model) fit else NULL
  )
}