fix interaction

Author: gravesti

R/trial_sequence.R

@@ -514,7 +514,12 @@ setMethod(
       period = as_formula(trial_period_terms, add = collection),
       stabilised = as_formula(get_stabilised_weights_terms(object), add = collection)
     )
-    adjustment <- unique(c(all.vars(formula_list$adjustment), all.vars(formula_list$stabilised)))
+    treatment <- all.vars(formula_list$treatment)
+    adjustment <- setdiff(
+      unique(c(all.vars(formula_list$adjustment), all.vars(formula_list$stabilised))),
+      treatment
+    )
+
     assert_names(
       adjustment,
       subset.of = colnames(object@data@data),
@@ -619,6 +624,7 @@ get_stabilised_weights_terms <- function(object) {
   stabilised_terms
 }
 
+# Update outcome model (formulas generally and variables in stabilised models)
 update_outcome_formula <- function(object) {
   assert_class(object, "trial_sequence")
 
@@ -638,7 +644,7 @@ update_outcome_formula <- function(object) {
   object@outcome_model@formula <- outcome_formula
 
   object@outcome_model@adjustment_vars <- unique(
-    c(all.vars(formula_list$adjustment), all.vars(formula_list$stabilised))
+    c(object@outcome_model@adjustment_vars, all.vars(formula_list$stabilised))
   )
 
   object

tests/testthat/test-trial_sequence.R

@@ -500,7 +500,29 @@ test_that("stabilised weight terms are included in outcome model", {
   )
 })
 
-# Expand
+
+test_that("interaction terms work as expected", {
+  result <- trial_sequence("PP") |>
+    set_data(data_censored) |>
+    set_outcome_model(adjustment_terms = ~ assigned_treatment * x2)
+
+  expect_equal(
+    result@outcome_model@formula,
+    outcome ~ assigned_treatment + x2 + followup_time + I(followup_time^2) +
+      trial_period + I(trial_period^2) + assigned_treatment:x2,
+    ignore_formula_env = TRUE
+  )
+  expect_equal(result@outcome_model@treatment_var, "assigned_treatment")
+  expect_equal(result@outcome_model@adjustment_vars, "x2") # shouldn't include treatment
+  expect_equal(
+    result@outcome_model@adjustment_terms, # can include treatment
+    ~ assigned_treatment * x2,
+    ignore_formula_env = TRUE
+  )
+})
+
+# Expand ---
+
 test_that("weights are 1 when not calculated by calculate_weights", {
   trial_ex <- TrialEmulation::trial_example