tsne-pca-umap-2d-3d-plot.py

import tensorflow as tf
import numpy as np

#For 3D Plot
import pandas as pd
import plotly.express as px

#For PCA
from sklearn.decomposition import PCA

#For t-SNE
from sklearn.manifold import TSNE

#For UMAP
import umap

#PARAMETERS
# PCA or TSNE or UMAP
dim_reduction_way = "TSNE"
#2D or 3D
plot_way = "3D"

# Prepare MNIST data.
from tensorflow.keras.datasets import mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = np.array(x_train, np.float32), np.array(x_test, np.float32)
x_train, x_test = x_train / 255., x_test / 255.

#reshape to (60000, 784) and (10000, 784)
x_train_shape = x_train.shape
x_train = x_train.reshape(x_train_shape[0], x_train_shape[1] * x_train_shape[2])
x_test_shape = x_test.shape
x_test = x_test.reshape(x_test_shape[0], x_test_shape[1] * x_test_shape[2])

#reduce to 1000
x_train = x_train[:1000]
y_train = y_train[:1000]
x_test = x_test[:1000]
y_test = y_test[:1000]


if dim_reduction_way == "PCA":
    x_train_transformed=PCA(n_components=8).fit_transform(x_train)
    x_test_transformed=PCA(n_components=8).fit_transform(x_test)

    #For 2D
    x_train_transformed_3d = x_train_transformed
    x_test_transformed_3d = x_test_transformed

    #For 3D (Same to 2d)
    x_train_transformed_2d = x_train_transformed
    x_test_transformed_2d = x_test_transformed
elif dim_reduction_way == "TSNE":
    #For 2D
    x_train_transformed_2d = TSNE(n_components=2).fit_transform(x_train)
    x_test_transformed_2d = TSNE(n_components=2).fit_transform(x_test)

    #For 3D
    x_train_transformed_3d = TSNE(n_components=3).fit_transform(x_train)
    x_test_transformed_3d = TSNE(n_components=3).fit_transform(x_test)
elif dim_reduction_way == "UMAP":
    #For 2D
    x_train_tranformed_2d = umap.UMAP(n_neighbors=5, n_components=2).fit(x_train)
    x_test_tranformed_2d = umap.UMAP(n_neighbors=5, n_components=2).fit(x_test)

    #For 3D
    x_train_tranformed_3d = umap.UMAP(n_neighbors=5, n_components=3).fit(x_train)
    x_test_tranformed_3d = umap.UMAP(n_neighbors=5, n_components=3).fit(x_test)


if plot_way == "2D":
    plt.scatter(x_train_transformed_2d[:,0],x_train_transformed_2d[:,1],s=20,c=y_train,cmap='tab10', marker="o", alpha=0.3)
    plt.scatter(x_test_transformed_2d[:,0],x_test_transformed_2d[:,1],s=20,c=y_test,cmap='tab10', marker="x", alpha=0.3)
    plt.colorbar()
elif plot_way == "3D":
    df_train=pd.DataFrame(x_train_transformed_3d[:, 0:3],columns=list("XYZ"))
    df_train["label"]=np.array(y_train)
    df_train['purpose'] = np.full(y_train.shape, "train")

    df_test=pd.DataFrame(x_test_transformed_3d[:, 0:3],columns=list("XYZ"))
    df_test["label"]=np.array(y_test)
    df_test['purpose'] = np.full(y_test.shape, "test")

    df_merged = pd.concat([df_train, df_test])
    fig = px.scatter_3d(df_merged, x='X', y='Y', z='Z',  color='label', symbol='purpose', size_max=4, opacity=0.7)
    fig.show()