R/fit_gwrf.R

In gwrf: Geographically Weighted Random Forests

#' Fit a Geographically Weighted Random Forest
#'
#' Fits a separate random-forest model for each focal observation or spatial
#' location using observations selected from a geographically defined local
#' neighborhood. Neighborhoods may be defined using individual data rows or
#' unique spatial locations, allowing the function to support repeated
#' observations at the same location, including spatial panel data.
#'
#' @param formula A model formula specifying the response and predictor
#'   variables.
#' @param data A data frame containing the response, predictors, and any
#'   location identifiers used in the model.
#' @param coords A numeric matrix or data frame with two columns containing
#'   the spatial coordinates associated with the rows of \code{data}.
#' @param bandwidth A positive numeric value defining the local neighborhood.
#'   When \code{adaptive = TRUE}, this is the number of neighboring rows or
#'   unique spatial locations included in each local neighborhood. When
#'   \code{adaptive = FALSE}, this is a fixed distance threshold expressed in
#'   the units of \code{coords}.
#' @param adaptive Logical. If \code{TRUE}, adaptive neighborhoods are defined
#'   using the nearest observations or unique locations. If \code{FALSE},
#'   fixed-distance neighborhoods are used.
#' @param kernel Character string specifying the spatial weighting kernel.
#'   The default is \code{"bisquare"}.
#' @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"}. All eligible observations associated
#'   with selected neighboring locations are retained for local model fitting.
#' @param num.trees Number of trees grown in each local random forest.
#' @param mtry Number of predictor variables randomly sampled as candidates at
#'   each split. If \code{NULL}, the value is determined by
#'   \code{ranger::ranger()}.
#' @param min.node.size Minimum terminal-node size used by each local random
#'   forest.
#' @param importance Character string specifying the variable-importance
#'   method passed to \code{ranger::ranger()}. The default is
#'   \code{"permutation"}.
#' @param use_case_weights Logical indicating whether spatial kernel weights
#'   are supplied to the local random forest as case weights.
#' @param focal_indices Optional integer vector identifying the focal
#'   observations for which local models should be fitted. If \code{NULL},
#'   local models are fitted for all eligible focal observations.
#' @param keep_local_models Logical indicating whether fitted local
#'   \code{ranger} model objects are retained in the returned object.
#' @param seed Optional integer random seed used for reproducible local
#'   random-forest fitting.
#' @param verbose Logical indicating whether progress messages are displayed
#'   during model fitting.
#'
#' @details
#' For each focal observation, the function constructs a spatial neighborhood,
#' fits a local random forest using the observations contained in that
#' neighborhood, and returns the focal prediction and predictor-importance
#' values. When \code{neighbor_unit = "location"}, adaptive bandwidth refers to
#' the number of nearest unique spatial locations rather than the number of
#' individual rows. This prevents repeated observations from the same location
#' from being treated as separate spatial neighbors.
#'
#' Spatial weights are determined by the selected kernel and the distances
#' between the focal location and neighboring observations or locations.
#' Variable importance describes predictive reliance within each fitted local
#' forest and does not indicate effect direction, statistical significance, or
#' causality.
#'
#' @return An object of class `"gwrf_fit"`. The object is a named list
#' containing the model call and specification, input data and coordinates,
#' neighborhood and random-forest settings, local model results, optional
#' fitted local models, and model diagnostics.
#'
#' The `local_results` component is a tibble with one row per fitted focal
#' observation and columns for the focal index, observed response, local
#' prediction, residual, local sample size, realized bandwidth, coordinates,
#' and, when available, local variable-importance values prefixed with
#' `"vi_"`.
#'
#' The `diagnostics` component is a list containing overall RMSE, MAE,
#' R-squared, and the number of focal models fitted.
#'
#' @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))
#'
#' fit <- fit_gwrf(
#'   y ~ x1 + x2,
#'   data = dat,
#'   coords = coords,
#'   bandwidth = 12,
#'   adaptive = TRUE,
#'   num.trees = 10,
#'   focal_indices = 1:3,
#'   seed = 1,
#'   verbose = FALSE
#' )
#'
#' fit
#' fit$local_results
#'
#' @seealso \code{\link[ranger]{ranger}}
#'
#' @export

fit_gwrf <- function(
    formula,
    data,
    coords,
    bandwidth,
    adaptive = TRUE,
    kernel = "bisquare",
    neighbor_unit = c("row", "location"),
    location_id = NULL,
    num.trees = 500,
    mtry = NULL,
    min.node.size = 5,
    importance = "permutation",
    use_case_weights = TRUE,
    focal_indices = NULL,
    keep_local_models = FALSE,
    seed = NULL,
    verbose = TRUE
) {
  check_gwrf_inputs(formula, data, coords)

neighbor_unit <- match.arg(neighbor_unit)

if (neighbor_unit == "location") {
  if (is.null(location_id)) {
    stop("location_id must be supplied when neighbor_unit = 'location'.")
  }

  if (length(location_id) == 1 && is.character(location_id) && location_id %in% names(data)) {
    location_id_vec <- data[[location_id]]
  } else {
    location_id_vec <- location_id
  }

  if (length(location_id_vec) != nrow(data)) {
    stop("location_id must be a column name in data or a vector with length nrow(data).")
  }

  if (any(is.na(location_id_vec))) {
    stop("location_id cannot contain NA values.")
  }
} else {
  location_id_vec <- NULL
}


  coords <- as.matrix(coords)

  if (is.null(focal_indices)) {
    focal_indices <- seq_len(nrow(data))
  }

  runner <- function(i) {
    fit_local_rf(
      data = data,
      formula = formula,
      coords = coords,
      focal_index = i,
      bandwidth = bandwidth,
      adaptive = adaptive,
      kernel = kernel,
      num.trees = num.trees,
      mtry = mtry,
      min.node.size = min.node.size,
      importance = importance,
      use_case_weights = use_case_weights,
      seed = seed,
      keep_model = keep_local_models,
      neighbor_unit = neighbor_unit,
      location_id = location_id_vec
    )
  }

  if (verbose) {
    local_fits <- pbapply::pblapply(focal_indices, runner)
  } else {
    local_fits <- lapply(focal_indices, runner)
  }

  local_tbl <- tibble::tibble(
    focal_index = vapply(local_fits, `[[`, integer(1), "focal_index"),
    observed = vapply(local_fits, `[[`, numeric(1), "observed"),
    prediction = vapply(local_fits, `[[`, numeric(1), "prediction"),
    residual = vapply(local_fits, `[[`, numeric(1), "residual"),
    n_local = vapply(local_fits, `[[`, numeric(1), "n_local"),
    local_bandwidth = vapply(local_fits, `[[`, numeric(1), "local_bandwidth"),
    x = coords[focal_indices, 1],
    y = coords[focal_indices, 2]
  )

  vi_list <- lapply(local_fits, `[[`, "variable_importance")
  vi_names <- unique(unlist(lapply(vi_list, names)))
  vi_names <- vi_names[!is.na(vi_names)]

  if (length(vi_names) > 0) {
    vi_df <- do.call(
      rbind,
      lapply(vi_list, function(v) {
        out <- rep(NA_real_, length(vi_names))
        names(out) <- vi_names
        if (!all(is.na(v))) {
          out[names(v)] <- unname(v)
        }
        out
      })
    )
    vi_df <- as.data.frame(vi_df)
    names(vi_df) <- paste0("vi_", names(vi_df))
    local_tbl <- dplyr::bind_cols(local_tbl, vi_df)
  }

  global_rmse <- sqrt(mean((local_tbl$observed - local_tbl$prediction)^2, na.rm = TRUE))
  global_mae <- mean(abs(local_tbl$observed - local_tbl$prediction), na.rm = TRUE)

  ss_res <- sum((local_tbl$observed - local_tbl$prediction)^2, na.rm = TRUE)
  ss_tot <- sum((local_tbl$observed - mean(local_tbl$observed, na.rm = TRUE))^2, na.rm = TRUE)
  global_r2 <- 1 - ss_res / ss_tot

  out <- list(
    call = match.call(),
    formula = formula,
    data = data,
    coords = coords,
    bandwidth = bandwidth,
    adaptive = adaptive,
    kernel = kernel,
    num.trees = num.trees,
    mtry = mtry,
    min.node.size = min.node.size,
    importance = importance,
    use_case_weights = use_case_weights,
    local_results = local_tbl,
    local_models = if (keep_local_models) local_fits else NULL,
    diagnostics = list(
      rmse = global_rmse,
      mae = global_mae,
      r2 = global_r2,
      n_focal = length(focal_indices)
    )
  )

  class(out) <- "gwrf_fit"
  out
}