diff --git a/R/02-pipeline.R b/R/02-pipeline.R index e7a350b..f277935 100644 --- a/R/02-pipeline.R +++ b/R/02-pipeline.R @@ -305,14 +305,14 @@ br_run <- function(obj, ..., group_by = NULL, run_parallel = 1L) { exponentiate <- dots[["exponentiate"]] if (is.null(group_by)) { - res <- runner(ms, obj@data, dots, obj@x, run_parallel) + res <- runner(ms, obj@data, dots, obj@y, obj@x, obj@x2, run_parallel) } else { obj@group_by <- group_by data_split <- obj@data |> named_group_split(obj@data[, group_by, drop = FALSE]) data_split[["All"]] <- obj@data res_list <- map(data_split, function(data) { - runner(ms, data, dots, obj@x, run_parallel) + runner(ms, data, dots, obj@y, obj@x, obj@x2, run_parallel, group_by) }) res <- list_transpose(res_list) res$models <- purrr::list_flatten(res$models) @@ -335,7 +335,11 @@ set_future_strategy <- function() { } } -runner <- function(ms, data, dots, x, run_parallel) { +runner <- function(ms, data, dots, y, x, x2, run_parallel, group_by = NULL) { + # Subset data to only necessary columns before model construction + necessary_cols <- get_necessary_columns(y, x, x2, group_by, colnames(data)) + data <- data[, necessary_cols, drop = FALSE] + log_n <- getOption("bregr.log_n", default = 100) assert_number_whole(log_n, min = 0, max = Inf, allow_infinite = TRUE) if (length(ms) >= log_n) { diff --git a/R/98-utils.R b/R/98-utils.R index ddd3871..7bfdf4a 100644 --- a/R/98-utils.R +++ b/R/98-utils.R @@ -24,8 +24,23 @@ get_vars <- function(text) { } } +get_necessary_columns <- function(y, x, x2, group_by, available_cols) { + # Get all variable names from y, x, x2, and group_by + necessary_vars <- merge_vars(y, x, x2, group_by) + + # Filter to only include columns that actually exist in the data + necessary_cols <- intersect(necessary_vars, available_cols) + + # Always include .row_names if it exists (added by tibble constructor) + if (".row_names" %in% available_cols) { + necessary_cols <- union(necessary_cols, ".row_names") + } + + return(necessary_cols) +} + merge_vars <- function(...) { - vars_list <- list(...) + vars_list <- list(...) |> unlist() rv <- NULL for (i in vars_list) { v <- unique(sapply(i, get_vars)) diff --git a/tests/testthat/test-optimization.R b/tests/testthat/test-optimization.R new file mode 100644 index 0000000..5ccf94f --- /dev/null +++ b/tests/testthat/test-optimization.R @@ -0,0 +1,103 @@ +test_that("unnecessary columns are removed from model objects", { + # Create a dataset with many unnecessary columns + set.seed(123) + n_rows <- 50 + n_extra_cols <- 20 + + # Create base data (what we actually need) + base_data <- data.frame( + y = rnorm(n_rows), + x1 = rnorm(n_rows), + x2 = rnorm(n_rows), + control1 = rnorm(n_rows) + ) + + # Add many unnecessary columns + extra_data <- replicate(n_extra_cols, rnorm(n_rows), simplify = FALSE) + names(extra_data) <- paste0("extra_col_", 1:n_extra_cols) + full_data <- cbind(base_data, extra_data) + + # Run modeling pipeline + result <- br_pipeline( + full_data, + y = "y", + x = c("x1", "x2"), + x2 = "control1", + method = "gaussian" + ) + + # Test that original data is preserved in breg object + expect_equal(ncol(result@data), ncol(full_data) + 1) # +1 for .row_names + + # Test that individual models only have necessary columns + model1 <- result@models[["x1"]] + model2 <- result@models[["x2"]] + + # Each model should only have y + focal_variable + control variables + # For x1: y, x1, control1 (3 columns) + expect_equal(ncol(model1$model), 3) + expect_equal(sort(colnames(model1$model)), sort(c("y", "x1", "control1"))) + + # For x2: y, x2, control1 (3 columns) + expect_equal(ncol(model2$model), 3) + expect_equal(sort(colnames(model2$model)), sort(c("y", "x2", "control1"))) + + # Test that results are still correct + manual_model1 <- lm(y ~ x1 + control1, data = base_data) + expect_equal(coef(manual_model1), coef(model1), tolerance = 1e-10) + + manual_model2 <- lm(y ~ x2 + control1, data = base_data) + expect_equal(coef(manual_model2), coef(model2), tolerance = 1e-10) +}) + +test_that("necessary columns are identified correctly", { + # Test the utility function directly + y <- c("response") + x <- c("focal1", "focal2", "poly(focal3, 2)") + x2 <- c("control1", "I(control2^2)") + group_by <- c("group_var") + available_cols <- c("response", "focal1", "focal2", "focal3", "control1", "control2", + "group_var", "extra1", "extra2", ".row_names") + + necessary <- get_necessary_columns(y, x, x2, group_by, available_cols) + + # Should include all variables referenced in y, x, x2, group_by + expected <- c("response", "focal1", "focal2", "focal3", "control1", "control2", + "group_var", ".row_names") + expect_setequal(necessary, expected) + + # Should not include extra columns + expect_false("extra1" %in% necessary) + expect_false("extra2" %in% necessary) +}) + +test_that("optimization works with group_by", { + set.seed(456) + n_rows <- 40 + + # Create test data with group variable + test_data <- data.frame( + y = rnorm(n_rows), + x1 = rnorm(n_rows), + control1 = rnorm(n_rows), + group_var = rep(c("A", "B"), each = n_rows/2), + extra1 = rnorm(n_rows), + extra2 = rnorm(n_rows), + extra3 = rnorm(n_rows) + ) + + # Run with group_by directly in pipeline + result <- breg(test_data) |> + br_set_y("y") |> + br_set_x("x1") |> + br_set_x2("control1") |> + br_set_model("gaussian") |> + br_run(group_by = "group_var") + + # Check that models only contain necessary columns + # Should have y, x1, control1 (3 columns) - group_var is used for splitting, not in model + for (model in result@models) { + expect_equal(ncol(model$model), 3) + expect_setequal(colnames(model$model), c("y", "x1", "control1")) + } +}) \ No newline at end of file