R/ransac_reg.R

In RANSAC: Robust Model Fitting Using the RANSAC Algorithm

#' Robust Linear Model Fitting via RANSAC
#'
#' Fits a robust linear model (`lm`) using the RANSAC algorithm.
#'
#' @param formula Model formula (as in `lm`).
#' @param data Data frame containing the model variables.
#' @param n_min Minimum number of points to fit the model (e.g., 2 for a straight line).
#' @param n_iter Number of iterations (higher values make the model more robust).
#' @param tol Absolute tolerance to consider a point as an inlier.
#' @param verbose If `TRUE`, shows progress messages.
#'
#' @return An `lm` model fitted only to the inliers, with an additional class `"ransac_model"` and an `"inliers"` attribute.
#'
#' @examples
#' set.seed(123)
#' x <- 1:10
#' y <- c(1, 2, 3, 4, 5, 6, 7, 8, 50, 60) # some outliers
#' data <- data.frame(x = x, y = y)
#' model <- ransac_reg(y ~ x, data = data, n_min = 2, tol = 5)
#' summary(model)
#'
#' @export
#' @importFrom stats lm predict
ransac_reg <- function(formula, data, n_min, n_iter = 100, tol = 0.2, verbose = FALSE) {
  best_inliers <- c()
  best_model <- NULL

  for (i in 1:n_iter) {
    sample_idx <- sample(1:nrow(data), n_min)
    sample_data <- data[sample_idx, ]

    model_try <- try(lm(formula, data = sample_data), silent = TRUE)
    if (inherits(model_try, "try-error")) next

    pred <- predict(model_try, newdata = data)
    y_var <- all.vars(formula)[1]
    error <- abs(data[[y_var]] - pred)
    inliers <- which(error < tol)

    if (length(inliers) > length(best_inliers)) {
      best_inliers <- inliers
      best_model <- model_try
      if (verbose) message("Iteration ", i, ": ", length(inliers), " inliers")
    }
  }

  final_model <- lm(formula, data = data[best_inliers, ])
  class(final_model) <- c("ransac_model", class(final_model))
  attr(final_model, "inliers") <- best_inliers
  return(final_model)
}