Give users the option to do this before or after the train/test split to prevent potential for data leakage (see https://arxiv.org/pdf/2207.07048.pdf).

{mice} package may be useful.