It seems there is a private method of UmapExperiment _prepare_umap_data who does the exact same things of the function prepare_umap_data in the umap_functions module, i.e. removing NaNs columns and train/test splitting.
Why code duplication? Which is the "most correct" one?