Changed tsne.py

Author: Nikita-Kedari

machine_learning/tsne.py

@@ -26,7 +26,6 @@
 """
 
 import doctest
-
 import numpy as np
 from sklearn.datasets import load_iris
 
@@ -38,49 +37,57 @@ def collect_dataset() -> tuple[np.ndarray, np.ndarray]:
     :return: Tuple containing feature matrix and target labels
 
     Example:
-    >>> x, y = collect_dataset()
-    >>> x.shape
+    >>> data_x, data_y = collect_dataset()
+    >>> data_x.shape
     (150, 4)
-    >>> y.shape
+    >>> data_y.shape
     (150,)
     """
     data = load_iris()
     return np.array(data.data), np.array(data.target)
 
 
-def compute_pairwise_affinities(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
+def compute_pairwise_affinities(data_x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
     """
     Computes pairwise affinities (P matrix) in high-dimensional space using Gaussian kernel.
 
-    :param x: Input data of shape (n_samples, n_features)
+    :param data_x: Input data of shape (n_samples, n_features)
     :param sigma: Variance (Bandwidth) of the Gaussian kernel
     :return: Symmetrized probability matrix p
 
     Example:
     >>> import numpy as np
-    >>> x = np.array([[0.0, 0.0], [1.0, 0.0]])
-    >>> p = compute_pairwise_affinities(x)
+    >>> data_x = np.array([[0.0, 0.0], [1.0, 0.0]])
+    >>> p = compute_pairwise_affinities(data_x)
     >>> float(round(p[0, 1], 3))
     0.25
     """
-    n_samples = x.shape[0]
-    sum_x = np.sum(np.square(x), axis=1)
-    d = np.add(np.add(-2 * np.dot(x, x.T), sum_x).T, sum_x)
+    n_samples = data_x.shape[0]
+    sum_x = np.sum(np.square(data_x), axis=1)
+    d = np.add(np.add(-2 * np.dot(data_x, data_x.T), sum_x).T, sum_x)
     p = np.exp(-d / (2 * sigma ** 2))
     np.fill_diagonal(p, 0)
     p /= np.sum(p)
     return (p + p.T) / (2 * n_samples)
 
 
-def compute_low_dim_affinities(y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+def compute_low_dim_affinities(embedding_y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     """
     Computes low-dimensional similarities (Q matrix) using Student-t distribution.
 
-    :param y: Low-dimensional embeddings (n_samples, n_components)
+    :param embedding_y: Low-dimensional embeddings (n_samples, n_components)
     :return: Tuple (q, num) where q is the probability matrix and num is numerator array
+
+    Example:
+    >>> embedding_y = np.array([[0.0, 0.0], [1.0, 0.0]])
+    >>> q, num = compute_low_dim_affinities(embedding_y)
+    >>> q.shape
+    (2, 2)
+    >>> num.shape
+    (2, 2)
     """
-    sum_y = np.sum(np.square(y), axis=1)
-    num = 1 / (1 + np.add(np.add(-2 * np.dot(y, y.T), sum_y).T, sum_y))
+    sum_y = np.sum(np.square(embedding_y), axis=1)
+    num = 1 / (1 + np.add(np.add(-2 * np.dot(embedding_y, embedding_y.T), sum_y).T, sum_y))
     np.fill_diagonal(num, 0)
     q = num / np.sum(num)
     return q, num
@@ -102,8 +109,8 @@ def apply_tsne(
     :return: Transformed dataset (low-dimensional embedding)
 
     Example:
-    >>> x, _ = collect_dataset()
-    >>> y_emb = apply_tsne(x, n_components=2, n_iter=50)
+    >>> data_x, _ = collect_dataset()
+    >>> y_emb = apply_tsne(data_x, n_components=2, n_iter=50)
     >>> y_emb.shape
     (150, 2)
     """
@@ -115,49 +122,54 @@ def apply_tsne(
     n_samples = data_x.shape[0]
 
     # Initialize low-dimensional map randomly
-    y = np.random.randn(n_samples, n_components) * 1e-4
+    y_emb = np.random.randn(n_samples, n_components) * 1e-4
     p = compute_pairwise_affinities(data_x)
     p = np.maximum(p, 1e-12)
 
     # Initialize parameters
-    y_inc = np.zeros_like(y)
+    y_inc = np.zeros_like(y_emb)
     momentum = 0.5
 
     for i in range(n_iter):
-        q, num = compute_low_dim_affinities(y)
+        q, num = compute_low_dim_affinities(y_emb)
         q = np.maximum(q, 1e-12)
 
         pq = p - q
 
         # Compute gradient
         d_y = 4 * (
-            np.dot((pq * num), y)
-            - np.multiply(np.sum(pq * num, axis=1)[:, np.newaxis], y)
+            np.dot((pq * num), y_emb)
+            - np.multiply(np.sum(pq * num, axis=1)[:, np.newaxis], y_emb)
         )
 
         # Update with momentum and learning rate
         y_inc = momentum * y_inc - learning_rate * d_y
-        y += y_inc
+        y_emb += y_inc
 
         # Adjust momentum halfway through
         if i == int(n_iter / 4):
             momentum = 0.8
 
-    return y
+    return y_emb
 
 
 def main() -> None:
     """
     Driver function for t-SNE demonstration.
+
+    Example:
+    >>> main()  # doctest: +ELLIPSIS
+    t-SNE embedding (first 5 points):
+    ...
     """
-    x, y_labels = collect_dataset()
-    y_emb = apply_tsne(x, n_components=2, n_iter=300)
+    data_x, data_y = collect_dataset()
+    y_emb = apply_tsne(data_x, n_components=2, n_iter=300)
     print("t-SNE embedding (first 5 points):")
     print(y_emb[:5])
 
     # Optional visualization (commented to avoid dependency)
     # import matplotlib.pyplot as plt
-    # plt.scatter(y_emb[:, 0], y_emb[:, 1], c=y_labels, cmap="viridis")
+    # plt.scatter(y_emb[:, 0], y_emb[:, 1], c=data_y, cmap="viridis")
     # plt.title("t-SNE Visualization of Iris Dataset")
     # plt.xlabel("Component 1")
     # plt.ylabel("Component 2")
@@ -181,13 +193,13 @@ def main() -> None:
 - n_iter: number of iterations for optimization
 
 Output:
-- y: numpy array of shape (n_samples, n_components)
+- y_emb: numpy array of shape (n_samples, n_components)
   Each row is the low-dimensional embedding of the corresponding high-dimensional point.
 
 How it works:
-1. Compute high-dimensional similarities (p matrix)
-2. Initialize low-dimensional map (y) randomly
-3. Compute low-dimensional similarities (q matrix)
-4. Minimize KL divergence between p and q using gradient descent
-5. Update y with momentum and learning rate iteratively
+1. Compute high-dimensional similarities (P matrix)
+2. Initialize low-dimensional map (y_emb) randomly
+3. Compute low-dimensional similarities (Q matrix)
+4. Minimize KL divergence between P and Q using gradient descent
+5. Update y_emb with momentum and learning rate iteratively
 """