Rotating 3D t-SNE animated gif scatterplot with matplotlib

So I’ve been through a few hours of searching + trial & error and came up with a simple solution to draw an animated GIF 3D scatterplot.
(minimum library installs, exclude bash commands)

X = data.iloc[:,0:-1]
Y = data.iloc[:,-1].astype('int')

from sklearn.manifold import TSNE
tsne = TSNE(n_components=3, random_state=RS, perplexity=10)
tsne_fit = tsne.fit_transform(X)

n_components should be set to 3 in order to draw a 3D plot.
perpliexity should be adjusted by trial and error to find the best value that represents your data. sqrt(N) is a good starting point.

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation
fig = plt.figure(figsize=(10,10))
ax = Axes3D(fig)
colors = 'b', 'r'
labels = 'Group1', 'Group2'

for i, c, label in zip(range(len(labels)), colors, labels):
    ax.scatter(tsne_fit[data['Group']==i, 0], tsne_fit[data['Group']==i, 1], tsne_fit[data['Group']==i, 2], s=30, c=c, label=label, alpha=0.5)
fig.legend()

Axes3D is for 3D plotting.
matplotlib.animation is for making animated GIF.

Draw the scatterplot. In my case, I used scatter() twice to label the outcome feature.
Added alpha=0.5 for better visualization when datapoints overlap.

def rotate(angle):
     ax.view_init(azim=angle)

angle = 3
ani = animation.FuncAnimation(fig, rotate, frames=np.arange(0, 360, angle), interval=50)
ani.save('inhadr_tsne1.gif', writer=animation.PillowWriter(fps=20))

Build an arbitrary function rotate() that updates the view of the plot. This function will be called by FuncAnimation().
The writer is set to PillowWriter since it’s included by default in matplotlib. But while searching, I found that in some cases there are some problems in the animation, and can be solved by using a different writer, such as FFMpegWriter.
angle=3 means the plot rotates 3 degrees every frame. (120 frames in total)
interval=50, fps=20 values can be tweaked to change the rotation speed of animation.

bingle bangle~

Took 10 seconds to draw this.
대만족! This is definitely going to be put in my next presentation!

CC BY-NC-SA 4.0 This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.

Leave a Comment

Time limit is exhausted. Please reload CAPTCHA.