feat: fixed number of steps

Author: aravindhebbali

R/ols-stepwise-hierarchical.R

@@ -13,27 +13,36 @@
 #' @param progress Logical; if \code{TRUE}, will display variable selection progress.
 #' @param details Logical; if \code{TRUE}, will print the regression result at
 #'   each step.
+#' @param steps Number of steps after which the stepwise procedures should stop.
 #'
 #' @examples
+#' # forward hierarchical selection
 #' model <- lm(y ~ ., data = surgical)
+#' ols_step_hierarchical(model)
+#'
+#' # backward hierarchical selection
 #' model <- lm(y ~ bcs + alc_heavy + pindex + enzyme_test + liver_test + alc_mod + age + gender, data = surgical)
+#' ols_step_hierarchical(model, forward = FALSE)
+#'
+#' # steps
 #' model <- lm(y ~ bcs + alc_heavy + pindex + enzyme_test + liver_test + age + gender + alc_mod, data = surgical)
+#' ols_step_hierarchical(model, steps = 2)
 #'
 #' @keywords internal
 #'
 #' @noRd
 #'
-ols_step_hierarchical <- function(model, p_value = 0.1, forward = TRUE, progress = FALSE, details = FALSE) {
+ols_step_hierarchical <- function(model, p_value = 0.1, forward = TRUE, progress = FALSE, details = FALSE, steps = NULL) {
 
   if (forward) {
-    ols_step_hierarchical_forward(model, p_value, progress, details)
+    ols_step_hierarchical_forward(model, p_value, progress, details, steps)
   } else {
-    ols_step_hierarchical_backward(model, p_value, progress, details)
+    ols_step_hierarchical_backward(model, p_value, progress, details, steps)
   }
 
 }
 
-ols_step_hierarchical_forward <- function(model, p_value = 0.1, progress = FALSE, details = FALSE) {
+ols_step_hierarchical_forward <- function(model, p_value = 0.1, progress = FALSE, details = FALSE, steps = NULL) {
 
   if (details) {
     progress <- FALSE
@@ -56,7 +65,7 @@ ols_step_hierarchical_forward <- function(model, p_value = 0.1, progress = FALSE
   if (progress || details) {
     ols_candidate_terms(nam, "forward")
   }
-    
+
   step     <- 0
   rsq      <- c()
   adjrsq   <- c()
@@ -68,7 +77,7 @@ ols_step_hierarchical_forward <- function(model, p_value = 0.1, progress = FALSE
 
   base_model <- lm(paste(response, "~", 1), data = l)
   rsq_base   <- summary(base_model)$r.squared
-    
+
   if (details) {
     ols_rsquared_init(NULL, "r2", response, rsq_base)
   }
@@ -77,6 +86,10 @@ ols_step_hierarchical_forward <- function(model, p_value = 0.1, progress = FALSE
     ols_progress_init("forward")
   }
 
+  if (!is.null(steps)) {
+    mlen_p <- steps
+  }
+
   for (i in seq_len(mlen_p)) {
     predictors  <- c(preds, all_pred[i])
     m           <- lm(paste(response, "~", paste(predictors, collapse = " + ")), l)
@@ -151,7 +164,7 @@ ols_step_hierarchical_forward <- function(model, p_value = 0.1, progress = FALSE
 
 }
 
-ols_step_hierarchical_backward <- function(model, p_value = 0.1, progress = FALSE, details = FALSE) {
+ols_step_hierarchical_backward <- function(model, p_value = 0.1, progress = FALSE, details = FALSE, steps = NULL) {
 
   if (details) {
     progress <- FALSE
@@ -166,9 +179,9 @@ ols_step_hierarchical_backward <- function(model, p_value = 0.1, progress = FALS
   cterms   <- nam
 
   if (progress || details) {
-    ols_candidate_terms(nam, "backward")  
+    ols_candidate_terms(nam, "backward")
   }
-    
+
   step     <- 0
   rsq      <- c()
   adjrsq   <- c()
@@ -179,7 +192,7 @@ ols_step_hierarchical_backward <- function(model, p_value = 0.1, progress = FALS
   rmse     <- c()
 
   rsq_base   <- summary(model)$r.squared
-    
+
   if (details) {
     ols_rsquared_init(NULL, "r2", response, rsq_base)
   }
@@ -194,7 +207,7 @@ ols_step_hierarchical_backward <- function(model, p_value = 0.1, progress = FALS
     m_sum       <- Anova(m)
     pvals       <- m_sum$`Pr(>F)`[1:i]
     p_vals      <- pvals[i]
-  
+
     if (details) {
       d <- data.frame(predictors = predictors, p_val = pvals)
       ols_stepwise_table_p(d, predictors, pvals)
@@ -222,6 +235,12 @@ ols_step_hierarchical_backward <- function(model, p_value = 0.1, progress = FALS
         ols_stepwise_details(step, preds, rpred, response, rsq1, "removed", "rsq")
       }
 
+      if (!is.null(steps)) {
+        if (step == steps) {
+          break
+        }
+      }
+
     } else {
 
       if (progress || details) {

tests/testthat/test-step-backward.R

@@ -29,6 +29,13 @@ test_that("backward hierarchical selection output matches the expected result",
   expect_equal(k$metrics$variable, c("alc_mod", "gender", "age", "liver_test"), ignore_attr = TRUE)
 })
 
+test_that("backward hierarchical selection stops after a specific number of steps", {
+  model <- lm(y ~ bcs + alc_heavy + pindex + enzyme_test + liver_test + age + gender + alc_mod, data = surgical)
+  k <- ols_step_backward_p(model, 0.1, hierarchical = TRUE, steps = 2)
+  expect_equal(k$metrics$step, 1:2)
+  expect_equal(k$metrics$variable, c("alc_mod", "gender"), ignore_attr = TRUE)
+})
+
 test_that("backward elimination output matches the expected result when steps in specified", {
   model <- lm(y ~ bcs + alc_heavy + pindex + enzyme_test + liver_test + age + gender + alc_mod, data = surgical)
   k <- ols_step_backward_p(model, steps = 2)

tests/testthat/test-step-forward.R

@@ -28,6 +28,13 @@ test_that("forward hierarchical selection output matches the expected result", {
   expect_equal(k$metrics$variable, c("bcs", "alc_heavy", "pindex", "enzyme_test"), ignore_attr = TRUE)
 })
 
+test_that("forward hierarchical selection stops after a specific number of steps", {
+  model <- lm(y ~ bcs + alc_heavy + pindex + enzyme_test + liver_test + age + gender + alc_mod, data = surgical)
+  k <- ols_step_forward_p(model, 0.1, hierarchical = TRUE, steps = 2)
+  expect_equal(k$metrics$step, 1:2)
+  expect_equal(k$metrics$variable, c("bcs", "alc_heavy"), ignore_attr = TRUE)
+})
+
 test_that("stepwise forward regression stops after a specific number of steps", {
   model <- lm(y ~ x1 + x2 + x3 + x4, data = cement)
   k <- ols_step_forward_p(model, steps = 2)