#%%
import numpy as np
import matplotlib.pyplot as plt

def fit_umap(X):
    from sklearn.manifold import UMAP
    umap = UMAP(n_neighbors=10, n_components=2, random_state=42)
    return umap.fit_transform(X)

X = np.random.rand(100, 2)
X_umap = fit_umap(X)

fig, ax = plt.subplots()
ax.scatter(X_umap[:, 0], X_umap[:, 1])
plt.show()
