1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
|
#!/usr/bin/env python
import matplotlib.pyplot as plt
import numba
import numpy as np
from mayavi import mlab
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import umap
digits = load_digits()
X_train, X_test, y_train, y_test = train_test_split(
digits.data, digits.target, stratify=digits.target, random_state=42
)
target_spaces = ["plane", "torus", "sphere"]
if "plane" in target_spaces:
# embed onto a plane
trans = umap.UMAP(
n_neighbors=10,
random_state=42,
metric="euclidean",
output_metric="euclidean",
init="spectral",
verbose=True,
).fit(X_train)
plt.scatter(
trans.embedding_[:, 0], trans.embedding_[:, 1], c=y_train, cmap="Spectral"
)
plt.show()
plt.close()
if "torus" in target_spaces:
# embed onto a torus
# note: this is a topological torus, not a geometric torus. Think
# Pacman, not donut.
@numba.njit(fastmath=True)
def torus_euclidean_grad(x, y, torus_dimensions=(2 * np.pi, 2 * np.pi)):
"""Standard euclidean distance.
..math::
D(x, y) = \sqrt{\sum_i (x_i - y_i)^2}
"""
distance_sqr = 0.0
g = np.zeros_like(x)
for i in range(x.shape[0]):
a = abs(x[i] - y[i])
if 2 * a < torus_dimensions[i]:
distance_sqr += a ** 2
g[i] = x[i] - y[i]
else:
distance_sqr += (torus_dimensions[i] - a) ** 2
g[i] = (x[i] - y[i]) * (a - torus_dimensions[i]) / a
distance = np.sqrt(distance_sqr)
return distance, g / (1e-6 + distance)
trans = umap.UMAP(
n_neighbors=10,
random_state=42,
metric="euclidean",
output_metric=torus_euclidean_grad,
init="spectral",
min_dist=0.15, # requires adjustment since the torus has limited space
verbose=True,
).fit(X_train)
mlab.clf()
x, y, z = np.mgrid[-3:3:50j, -3:3:50j, -3:3:50j]
# Plot a torus
R = 2
r = 1
values = (R - np.sqrt(x ** 2 + y ** 2)) ** 2 + z ** 2 - r ** 2
mlab.contour3d(x, y, z, values, color=(1.0, 1.0, 1.0), contours=[0])
# torus angles -> 3D
x = (R + r * np.cos(trans.embedding_[:, 0])) * np.cos(trans.embedding_[:, 1])
y = (R + r * np.cos(trans.embedding_[:, 0])) * np.sin(trans.embedding_[:, 1])
z = r * np.sin(trans.embedding_[:, 0])
pts = mlab.points3d(
x, y, z, y_train, colormap="spectral", scale_mode="none", scale_factor=0.1
)
mlab.savefig('torus.png')
if "sphere" in target_spaces:
# embed onto a sphere
trans = umap.UMAP(
n_neighbors=10,
random_state=42,
metric="euclidean",
output_metric="haversine",
init="spectral",
min_dist=0.15, # requires adjustment since the sphere has limited space
verbose=True,
).fit(X_train)
mlab.clf()
x, y, z = np.mgrid[-3:3:50j, -3:3:50j, -3:3:50j]
# Plot a sphere
r = 3
values = x ** 2 + y ** 2 + z ** 2 - r ** 2
mlab.contour3d(x, y, z, values, color=(1.0, 1.0, 1.0), contours=[0])
# latitude, longitude -> 3D
x = r * np.sin(trans.embedding_[:, 0]) * np.cos(trans.embedding_[:, 1])
y = r * np.sin(trans.embedding_[:, 0]) * np.sin(trans.embedding_[:, 1])
z = r * np.cos(trans.embedding_[:, 0])
pts = mlab.points3d(
x, y, z, y_train, colormap="spectral", scale_mode="none", scale_factor=0.2
)
mlab.savefig('sphere.png')
|