`step_linadjust` for linear adjustment of covariates

[mattansb]

A common step in data processing - especially for clustering - is batch effect removal: in general, this is (linear) adjustment of covariates, typically to remove the "effect" of some categorical variable.

This is implemented in the {limma} package - here is an example for removing differences between species in flipper_len and body_mass to better help in detecting six differences:

# pak::pak("limma")

library(ggplot2)
library(patchwork)

data(penguins)
penguins <- na.omit(penguins)

penguins_size <- penguins |>
  subset(select = c(flipper_len, body_mass))

penguins_size_adj <- limma::removeBatchEffect(
  x = t(penguins_size),
  batch = penguins$species,
  design = model.matrix(~sex, data = penguins)
) |>
  t() |>
  as.data.frame()


p <- ggplot(
  penguins,
  aes(x = flipper_len, y = body_mass, color = penguins$sex)
) +
  geom_point() +
  stat_ellipse() +
  labs(
    x = "Flipper Length (mm)",
    y = "Body Mass (g)",
    color = "Island"
  ) +
  theme_minimal()

p + (p + penguins_size_adj + ggtitle("limma::removeBatchEffect of Species"))

Here is the same implementation using {recipes}:

step_linadjust

# 1. User-facing function
step_linadjust <- function(
  recipe,
  ...,
  role = NA,
  trained = FALSE,
  remove_vars = NULL, # The "Batch" to remove
  keep_vars = NULL, # The "Design" to preserve
  models = NULL,
  drop = c("both", "remove", "none"), # drop the nuisance variables from the output?
  skip = FALSE,
  id = rand_id("linadjust")
) {
  add_step(
    recipe,
    step_linadjust_new(
      terms = enquos(...),
      role = role,
      trained = trained,
      remove_vars = remove_vars,
      keep_vars = keep_vars,
      models = models,
      drop = drop,
      skip = skip,
      id = id
    )
  )
}

# 2. Step Constructor
step_linadjust_new <- function(
  terms,
  role,
  trained,
  remove_vars,
  keep_vars,
  models,
  drop,
  skip,
  id
) {
  step(
    subclass = "linadjust",
    terms = terms,
    role = role,
    trained = trained,
    remove_vars = remove_vars,
    keep_vars = keep_vars,
    models = models,
    drop = drop,
    skip = skip,
    id = id
  )
}

# 3. Prep Method (Fits the model Y ~ Remove + Keep)
prep.step_linadjust <- function(x, training, info = NULL) {
  # Identify target columns (Y)

  col_names <- recipes_eval_select(x$terms, training, info)

  # Identify Nuisance columns (Batch)
  if (is.null(x$remove_vars)) {
    cli::cli_abort(
      c(
        "The `remove_vars` argument must be specified.",
        "i" = "This is the variable(s) you want to remove the effect of."
      )
    )
  }

  remove_names <- recipes_eval_select(x$remove_vars, training, info)

  # Identify Preserved columns (Design)
  # Handle case where keep_vars is NULL
  if (!is.null(x$keep_vars)) {
    keep_names <- recipes_eval_select(x$keep_vars, training, info)

    if (any(keep_names %in% remove_names)) {
      cli::cli_abort(
        c(
          "The `keep_vars` and `remove_vars` selectors must be disjoint.",
          "x" = "The following variables are in both: {intersect(keep_names, remove_names)}"
        )
      )
    }
  } else {
    keep_names <- NULL
  }

  for (ic in c(remove_names, keep_names)) {
    if (is.factor(training[[ic]])) {
      training[[ic]] <- C(droplevels(training[[ic]]), contr.sum)
    } else if (is.numeric(training[[ic]])) {
      training[[ic]] <- scale(training[[ic]])
    } else {
      cli::cli_abort(
        c(
          "The `remove_vars` and `keep_vars` selectors must be either factors or numeric.",
          "x" = "The following variable is neither: {ic}"
        )
      )
    }
  }

  model_list <- list()

  for (col in col_names) {
    # Create formula: Target ~ Remove1 + Keep1 + ...
    # We combine both sets of variables for the fit
    ff <- reformulate(
      response = col,
      termlabels = c(remove_names, keep_names)
    )

    # Fit and store the model
    model_list[[col]] <- butcher::butcher(lm(ff, data = training))
  }

  drop <- match.arg(x$drop, choices = c("both", "remove", "none"))

  step_linadjust_new(
    terms = col_names,
    role = x$role,
    trained = TRUE,
    remove_vars = remove_names,
    keep_vars = keep_names,
    models = model_list,
    drop = drop,
    skip = x$skip,
    id = x$id
  )
}

# 4. Bake Method (Subtracts ONLY the Remove effect)
bake.step_linadjust <- function(object, new_data, ...) {
  # Get names of the variables we want to remove effects for
  # We need to re-evaluate the selector to get string names
  remove_names <- names(object$remove_vars)
  keep_names <- names(object$keep_vars)

  for (col in names(object$models)) {
    model <- object$models[[col]]

    # Crucial Step: use type = "terms"
    # This returns a matrix with one column per independent variable,
    # representing that variable's contribution to the prediction.
    # It handles factors (dummification) automatically.
    term_preds <- predict(model, newdata = new_data, type = "terms")

    # Identify which columns in the term matrix correspond to our `remove_vars`
    # Note: `predict` names columns by the variable name.
    cols_to_subtract <- intersect(colnames(term_preds), remove_names)

    if (length(cols_to_subtract) > 0) {
      # Sum the effects of the nuisance variables
      nuisance_effect <- rowSums(term_preds[, cols_to_subtract, drop = FALSE])

      # Subtract nuisance effect from original data
      # Result = (Signal + Batch + Noise) - (Batch) = Signal + Noise
      new_data[[col]] <- new_data[[col]] - nuisance_effect
    }
  }

  if (object$drop == "remove") {
    new_data <- new_data[,
      !(colnames(new_data) %in% remove_names),
      drop = FALSE
    ]
  } else if (object$drop == "both") {
    new_data <- new_data[,
      !(colnames(new_data) %in% c(remove_names, keep_names)),
      drop = FALSE
    ]
  }

  tibble::as_tibble(new_data)
}

library(recipes)
#> Loading required package: dplyr
#> 
#> Attaching package: 'dplyr'
#> The following objects are masked from 'package:stats':
#> 
#>     filter, lag
#> The following objects are masked from 'package:base':
#> 
#>     intersect, setdiff, setequal, union
#> 
#> Attaching package: 'recipes'
#> The following object is masked from 'package:stats':
#> 
#>     step

rec <- recipe(~ flipper_len + body_mass + species + sex, data = penguins) |>
  step_linadjust(
    flipper_len,
    body_mass,
    remove_vars = vars(species),
    keep_vars = vars(sex)
  )


p +
  (p +
    prep(rec) |> bake(new_data = penguins) +
    ggtitle("step_linadjust Species"))
#> Registered S3 method overwritten by 'butcher':
#>   method                 from    
#>   as.character.dev_topic generics

^{Created on 2026-07-02 with reprex v2.1.1}

[mattansb]

Not sure if this falls under the scope of {embed}?

[topepo]

Yeah, probably embed. Bioconductor dependencies are possible but have high maintenance. IIRC, this method just does some linear regression work with covariates. Can we replicate this natively instead of using the dependency?

[mattansb]

Yes, if you look at the folded code above you will see I've completely implanted the functionality in tidymodels-ese (:

[topepo]

Ok, I'll move this to embed. Can you make a PR? We have a developer skill if you are using an LLM to assist you as well as a page with guidelines and notes.

Also, would you be open to changing the name? Maybe step_adjust_batches_linear or something similar? That's long but self-describing and the last term _linear could be replaced if similar steps were added.

[EmilHvitfeldt]

and please feel free ot ask any questions along the way, i can help guide if needed