UMAP results not robust to StratifiedCV

Nobody has claimed this yet.

This is something I've been able to reproduce over 3 different datasets, although all 3 are text embeddings. Each case is a binary classification problem where the positive class frequency is roughly 1% and the dataset size is ~ 3-5k. There are ~ 500 features. I've tried multiple parameter combinations for n_neighbors, target_weight, and min_dist. All plots below use the following parameters: (n_neighbors=50, metric='cosine', random_state=42, target_weight=0.9)

If I take all my labeled data and run UMAP in supervised mode I get a nice separation between my positive and negative classes like so:

However, if I hold out 10% of the data using sklearn's StratifiedShuffleSplit and set those labels to -1, the resulting plots look very different, even if I set n_neighbors to 100:

This makes me much less confident about using UMAP as a metric learning tool. Any thoughts as to why this is happening (or more likely what I'm doing wrong) would be greatly appreciated.

No source file or test is named. Start by reproducing the supervised UMAP runs with the stated n_neighbors, metric, random_state, target_weight, and StratifiedShuffleSplit setup, then inspect how labels set to -1 are handled. Done means determining whether the changed embedding is expected or a bug, with the behavior and reproduction covered by an appropriate test or documentation.

Written by the indexing model from the issue text.