plot_umap.py

import umap
import matplotlib.pyplot as plt

from mpl_toolkits.mplot3d import axes3d

def draw_umap(data, n_neighbors=15, min_dist=0.1, n_components=2, metric='euclidean', title='', color=None):
    fit = umap.UMAP(
        n_neighbors=n_neighbors,
        min_dist=min_dist,
        n_components=n_components,
        metric=metric
    )
    u = fit.fit_transform(data);
    fig = plt.figure()
    
#    if not color:
#        color = preprocessing.MinMaxScaler().fit_transform(data[:,:4])
        
    if n_components == 1:
        ax = fig.add_subplot(111)
        ax.scatter(u[:,0], range(len(u)), c=color)
    if n_components == 2:
        ax = fig.add_subplot(111)
        ax.scatter(u[:,0], u[:,1], c=color)
    if n_components == 3:
        ax = fig.add_subplot(111, projection='3d')
        ax.scatter(u[:,0], u[:,1], u[:,2], c=color)
    plt.title(title, fontsize=18)