t-SNE: Split out affinity, initialization and object preparation

Author: pavlin-policar

Orange/projection/manifold.py

@@ -400,7 +400,7 @@ def __init__(self, n_components=2, perplexity=30, learning_rate=200,
         self.callbacks_every_iters = callbacks_every_iters
         self.random_state = random_state
 
-    def fit(self, X: np.ndarray, Y: np.ndarray = None) -> openTSNE.TSNEEmbedding:
+    def compute_affinities(self, X):
         # Sparse data are not supported
         if sp.issparse(X):
             raise TypeError(
@@ -415,41 +415,75 @@ def fit(self, X: np.ndarray, Y: np.ndarray = None) -> openTSNE.TSNEEmbedding:
             if not isinstance(self.perplexity, Iterable):
                 raise ValueError(
                     "Perplexity should be an instance of `Iterable`, `%s` "
-                    "given." % type(self.perplexity).__name__)
+                    "given." % type(self.perplexity).__name__
+                )
             affinities = openTSNE.affinity.Multiscale(
-                X, perplexities=self.perplexity, metric=self.metric,
-                method=self.neighbors, random_state=self.random_state, n_jobs=self.n_jobs)
+                X,
+                perplexities=self.perplexity,
+                metric=self.metric,
+                method=self.neighbors,
+                random_state=self.random_state,
+                n_jobs=self.n_jobs,
+            )
         else:
             if isinstance(self.perplexity, Iterable):
                 raise ValueError(
                     "Perplexity should be an instance of `float`, `%s` "
-                    "given." % type(self.perplexity).__name__)
+                    "given." % type(self.perplexity).__name__
+                )
             affinities = openTSNE.affinity.PerplexityBasedNN(
-                X, perplexity=self.perplexity, metric=self.metric,
-                method=self.neighbors, random_state=self.random_state, n_jobs=self.n_jobs)
+                X,
+                perplexity=self.perplexity,
+                metric=self.metric,
+                method=self.neighbors,
+                random_state=self.random_state,
+                n_jobs=self.n_jobs,
+            )
 
-        # Create an initial embedding
+        return affinities
+
+    def compute_initialization(self, X):
+        # Compute the initial positions of individual points
         if isinstance(self.initialization, np.ndarray):
             initialization = self.initialization
         elif self.initialization == "pca":
             initialization = openTSNE.initialization.pca(
-                X, self.n_components, random_state=self.random_state)
+                X, self.n_components, random_state=self.random_state
+            )
         elif self.initialization == "random":
             initialization = openTSNE.initialization.random(
-                X, self.n_components, random_state=self.random_state)
+                X, self.n_components, random_state=self.random_state
+            )
         else:
             raise ValueError(
                 "Invalid initialization `%s`. Please use either `pca` or "
-                "`random` or provide a numpy array." % self.initialization)
+                "`random` or provide a numpy array." % self.initialization
+            )
 
-        embedding = openTSNE.TSNEEmbedding(
-            initialization, affinities, learning_rate=self.learning_rate,
-            theta=self.theta, min_num_intervals=self.min_num_intervals,
-            ints_in_interval=self.ints_in_interval, n_jobs=self.n_jobs,
+        return initialization
+
+    def prepare_embedding(self, affinities, initialization):
+        """Prepare an embedding object with appropriate parameters, given some
+        affinities and initialization."""
+        return openTSNE.TSNEEmbedding(
+            initialization,
+            affinities,
+            learning_rate=self.learning_rate,
+            theta=self.theta,
+            min_num_intervals=self.min_num_intervals,
+            ints_in_interval=self.ints_in_interval,
+            n_jobs=self.n_jobs,
             negative_gradient_method=self.negative_gradient_method,
-            callbacks=self.callbacks, callbacks_every_iters=self.callbacks_every_iters,
+            callbacks=self.callbacks,
+            callbacks_every_iters=self.callbacks_every_iters,
         )
 
+    def fit(self, X: np.ndarray, Y: np.ndarray = None) -> openTSNE.TSNEEmbedding:
+        # Compute affinities and initial positions and prepare the embedding object
+        affinities = self.compute_affinities(X)
+        initialization = self.compute_initialization(X)
+        embedding = self.prepare_embedding(affinities, initialization)
+
         # Run standard t-SNE optimization
         embedding.optimize(
             n_iter=self.early_exaggeration_iter, exaggeration=self.early_exaggeration,