import numpy as np
import plotly.graph_objects as go
from scipy.spatial.distance import pdist, squareform
# --- 1. CONFIGURATION & DATA GENERATION ---
np.random.seed(13)
n_points = 10
n_iter = 500
learning_rate = 10
momentum = 0.75
# Generate 3D data: 3 clusters
# We shift the Z-axis by +5 so the original data floats above the Z=0 projection plane
cluster_1 = np.random.normal(loc=[0, 0, 10], scale=1.0, size=(4, 3))
cluster_2 = np.random.normal(loc=[30, 10, 8], scale=2.0, size=(3, 3))
cluster_3 = np.random.normal(loc=[-30, -10, 6], scale=3.0, size=(3, 3))
X_3d = np.vstack([cluster_1, cluster_2, cluster_3])
colors = ['red'] * 4 + ['green'] * 3 + ['blue'] * 3
# --- 2. MINI T-SNE IMPLEMENTATION ---
def compute_joint_probabilities(X, sigma=1.0):
dist = squareform(pdist(X, 'sqeuclidean'))
P = np.exp(-dist / (2 * sigma**2))
np.fill_diagonal(P, 0)
P = P + P.T
P = P / np.sum(P)
return P
def compute_gradient(P, Y):
n = len(Y)
dist_Y = squareform(pdist(Y, 'sqeuclidean'))
inv_dist = 1.0 / (1.0 + dist_Y)
np.fill_diagonal(inv_dist, 0)
Q = inv_dist / np.sum(inv_dist)
grads = np.zeros_like(Y)
PQ_diff = (P - Q) * inv_dist
for i in range(n):
diff = Y[i] - Y
grads[i] = 4 * np.sum((PQ_diff[i, :][:, np.newaxis] * diff), axis=0)
return grads
# Initialize 2D embedding (small random numbers)
Y = np.random.normal(0, 0.0001, size=(n_points, 2))
Y_history = [Y.copy()]
velocity = np.zeros_like(Y)
# Run Optimization
P = compute_joint_probabilities(X_3d)
for i in range(n_iter):
grads = compute_gradient(P, Y)
velocity = momentum * velocity - learning_rate * grads
Y = Y + velocity
Y_history.append(Y.copy())
# --- 3. CALCULATE FIXED RANGES ---
# We need to know the max extent of the data AND the final embedding
# to set the axis ranges so the camera doesn't jump.
final_Y = Y_history[-1]
all_x = np.concatenate([X_3d[:, 0], final_Y[:, 0]])
all_y = np.concatenate([X_3d[:, 1], final_Y[:, 1]])
all_z = np.concatenate([X_3d[:, 2], np.zeros(n_points)]) # Projections are at z=0
# Add a little padding
x_range = [np.min(all_x) - 2, np.max(all_x) + 2]
y_range = [np.min(all_y) - 2, np.max(all_y) + 2]
z_range = [0, np.max(all_z) + 1] # From floor to top of data
# --- 4. PLOTLY VISUALIZATION ---
fig = go.Figure()
# -- Trace 1: Original 3D Data (Static) --
fig.add_trace(go.Scatter3d(
x=X_3d[:, 0], y=X_3d[:, 1], z=X_3d[:, 2],
mode='markers',
marker=dict(size=6, color=colors, opacity=0.5), # Slightly transparent
name='Original Data'
))
# -- Trace 2: Projected Points (Dynamic) --
# Initial state (Step 0)
fig.add_trace(go.Scatter3d(
x=Y_history[0][:, 0],
y=Y_history[0][:, 1],
z=np.zeros(n_points), # Fixed at Z=0
mode='markers',
marker=dict(size=10, color=colors, symbol='circle', opacity=1.0),
name='t-SNE Projection'
))
# -- Frames for Animation --
frames = []
for k, Y_step in enumerate(Y_history[1::5]):
frames.append(go.Frame(
data=[
go.Scatter3d(), # Keep Trace 1 static
go.Scatter3d( # Update Trace 2
x=Y_step[:, 0],
y=Y_step[:, 1],
z=np.zeros(n_points)
)
],
name=str(k*5)
))
fig.frames = frames
# -- Layout Configuration --
fig.update_layout(
title="3D Data and 2D t-SNE Projection",
width=1000, height=500,
scene=dict(
xaxis=dict(range=x_range, title="X"),
yaxis=dict(range=y_range, title="Y"),
zaxis=dict(range=z_range, title="Z"),
aspectmode='manual',
aspectratio=dict(x=1.5, y=1.5, z=1.0) # Flatten the Z visual slightly
),
updatemenus=[dict(
type="buttons",
buttons=[dict(label="Play",
method="animate",
args=[None, dict(frame=dict(duration=20, redraw=True), fromcurrent=True)])],
x=0.1, y=0.05
)],
sliders=[dict(
steps=[dict(method='animate',
args=[[str(k*5)], dict(mode='immediate', frame=dict(duration=0, redraw=True))],
label=str(k*5)) for k in range(len(Y_history[1::5]))],
currentvalue=dict(prefix="Iteration: "),
x=0.1, y=0, len=0.9
)]
)
fig.show()
