MICE on train/test split for machine learning prediction

[AlirezaShiraniBidabadi]

hello everyone i am using mice to generate five imputed datasets (five train and five test) then i would use an AI model for each data and finally using avg( based on rubin's rule) i will combine the predictions. now my question is how should i treat the target column. if i keep it for train data and make it null for test data then there is an inconsistency since when the imputed model look at null in test data it would probably return null and therefore the test data remains not imputed(this is my case) but in the document i have seen that it is recommended to keep the target col in train data. now i am also thinking of filling the target col for test with noise(random selection of train labels) i could not find any specific doc on this i found that you should use ignore to make sure test data is not present during imputation but how about the target col? here is my code ''
library(mice)
library(caret)
library(haven)

continuous_vars <- c("features")
binary_vars <- c("features"")
categorical_vars <- c("features")

Fix types

df[continuous_vars] <- lapply(df[continuous_vars], as.numeric)
df[binary_vars] <- lapply(df[binary_vars], as.factor)
df[categorical_vars] <- lapply(df[categorical_vars], as.factor)
df$Target<- as.factor(df$Target)

set.seed(123)
....
....
....
...code for splitting....
train_data <- df[ train_idx, ]
test_data <- df[-train_idx, ]

train_target <- train_data$Target
test_target <- test_data$Target

train_data$Target <- train_target
test_data$Target <- NA

── 3. Stack train + test; mark test rows with ignore = TRUE ───────────────

n_train <- nrow(train_data)
n_test <- nrow(test_data)

combined <- rbind(train_data, test_data)
ignore_flags <- c(rep(FALSE, n_train),
rep(TRUE, n_test))

── 4. Configure imputation methods ───────────────────────────────────────

ini <- mice(combined, maxit = 0, printFlag = FALSE)
meth <- ini$method

meth[continuous_vars] <- "pmm"
meth[binary_vars] <- "logreg"
meth[categorical_vars] <- "polyreg"
meth["Target"] <- ""

── 5. Run MICE ────────────────────────────────────────────────────────────

cat("-> Running MICE (Target as predictor in train, absent in test)...\n")
mice_fit <- mice(
combined,
m = 5,
maxit = 10,
method = meth,
ignore = ignore_flags,
printFlag = FALSE,
seed = 123
)

── 6. Extract train and test using mice filter() ─────────────────────────

mice_train <- filter(mice_fit, !ignore_flags)
mice_test <- filter(mice_fit, ignore_flags)

for (m_idx in 1:5) {

final_train <- complete(mice_train, m_idx)
final_test <- complete(mice_test, m_idx)

rownames(final_train) <- NULL
rownames(final_test) <- NULL

# Drop the NA Target that came through, reattach true labels

final_train$Target<- train_target
final_test$Target<- test_target

final_train$Set_Type <- "Train"
final_test$Set_Type <- "Test"

...saving files ''''' but in my case test data stays unimputed.

[ahmedTmersal]

Your approach is correct — the ignore = TRUE flag is exactly right. The issue is with this line:
rmeth["Target"] <- ""

This tells MICE to not impute Target at all, so test rows with NA Target stay NA. That's actually fine and expected — you don't want to impute the target.
Your real problem: Test features aren't being imputed either?
Check by printing:
rcomplete(mice_test, 1)[1:5, ]
If features are imputed but Target is NA — that's correct behavior. You then replace it with:
rfinal_test$Target <- test_target # ← you're already doing this!
So your code is actually fine. The "test data stays unimputed" might just be the Target column appearing as NA before line final_test$Target <- test_target runs.
Bottom line: Don't impute the target. Your ignore flag handles feature imputation correctly, and you restore true labels afterward. This is the right approach per Rubin's rules.

[AlirezaShiraniBidabadi]

thanks for your reply....actually yes when i exclude the target column my test features are imputed as well as the train data but when i set test data's target col as null it does not impute features in the test data. i think i should use random noise instead

[ahmedTmersal]

Before switching to random noise, can you check the predictorMatrix?
rini$predictorMatrix["continuous_var_example", "Target"]

If Target is being used as a predictor for your features (value = 1), then setting it to NA for all test rows means MICE has no information from Target to impute those test rows — which would explain why features stay unimputed only when Target is NA.
Two options instead of noise:

Remove Target as a predictor for other variables (set its column to 0 in predictorMatrix) — cleaner, avoids leaking noisy info.
Use mean/mode imputation placeholder for Target in test (not random noise) just to give MICE something to condition on, then drop it and replace with true labels after.

Random noise could introduce bias since you're injecting fake correlations between Target and other features. Removing Target from the predictor matrix for feature-imputation is safer and more standard practice.

[AlirezaShiraniBidabadi]

hi again thank you the problem was solved with the correct use of predictorMatrix

[thomvolker]

Hi @AlirezaShiraniBidabadi Thanks for your interest in using the mice package. We are currently finalizing a preprint about this problem. I'm assuming the target is completely observed in your training data, and here only the predictors need to be imputed. In the test data, predictors can be missing and the target is obviously not available. The solution is as follows. You can marginalize over the distribution of the outcome variable. In other words, you treat the outcome in the test set as a missing value that you impute. These imputations serve solely as a vehicle to impute the missing test predictors, and can be discarded afterwards. In this way, you can use the relation between the target variable and the missing data to create more accurate imputations, resulting in a better prediction model and ultimately in better predictions.

Hope this helps!

[AlirezaShiraniBidabadi]

Hi @thomvolker thank you for your detailed answer