Merge pull request biolab#3592 from pavlin-policar/tsne-faster

[FIX] t-SNE speed-ups

Author: lanzagar

Orange/widgets/unsupervised/owtsne.py

@@ -18,10 +18,11 @@
 
 
 class TSNERunner:
-    def __init__(self, tsne: TSNEModel, step_size=50):
+    def __init__(self, tsne: TSNEModel, step_size=50, exaggeration=1):
         self.embedding = tsne
         self.iterations_done = 0
         self.step_size = step_size
+        self.exaggeration = exaggeration
 
         # Larger data sets need a larger number of iterations
         if self.n_samples > 100_000:
@@ -43,7 +44,7 @@ def run_optimization(self):
         while not current_iter >= total_iterations:
             # Switch to normal regime if early exaggeration phase is over
             if current_iter >= self.early_exagg_iter:
-                exaggeration, momentum = 1, 0.8
+                exaggeration, momentum = self.exaggeration, 0.8
 
             # Resume optimization for some number of steps
             self.embedding.optimize(
@@ -73,7 +74,7 @@ class OWtSNE(OWDataProjectionWidget):
     settings_version = 3
     max_iter = Setting(300)
     perplexity = Setting(30)
-    multiscale = Setting(True)
+    multiscale = Setting(False)
     exaggeration = Setting(1)
     pca_components = Setting(20)
     normalize = Setting(True)
@@ -110,6 +111,13 @@ def __init__(self):
         self.__in_next_step = False
         self.__draw_similar_pairs = False
 
+        def reset_needs_to_draw():
+            self.needs_to_draw = True
+
+        self.needs_to_draw = True
+        self.__timer_draw = QTimer(self, interval=2000,
+                                   timeout=reset_needs_to_draw)
+
     def _add_controls(self):
         self._add_controls_start_box()
         super()._add_controls()
@@ -258,6 +266,8 @@ def pca_preprocessing(self):
     def __start(self):
         self.pca_preprocessing()
 
+        self.needs_to_draw = True
+
         # We call PCA through fastTSNE because it involves scaling. Instead of
         # worrying about this ourselves, we'll let the library worry for us.
         initialization = TSNE.default_initialization(
@@ -281,10 +291,13 @@ def __start(self):
             n_components=2, perplexity=perplexity, multiscale=self.multiscale,
             early_exaggeration_iter=0, n_iter=0, initialization=initialization,
             exaggeration=self.exaggeration, neighbors=neighbor_method,
-            negative_gradient_method=gradient_method, random_state=0
+            negative_gradient_method=gradient_method, random_state=0,
+            theta=0.8,
         )(self.pca_data)
 
-        self.tsne_runner = TSNERunner(self.projection, step_size=50)
+        self.tsne_runner = TSNERunner(
+            self.projection, step_size=20, exaggeration=self.exaggeration
+        )
         self.tsne_iterator = self.tsne_runner.run_optimization()
         self.__set_update_loop(self.tsne_iterator)
         self.progressBarInit(processEvents=None)
@@ -305,6 +318,7 @@ def __set_update_loop(self, loop):
             self.runbutton.setText("Stop")
             self.__state = OWtSNE.Running
             self.__timer.start()
+            self.__timer_draw.start()
         else:
             self.setBlocking(False)
             self.setStatusMessage("")
@@ -313,6 +327,7 @@ def __set_update_loop(self, loop):
             if self.__state == OWtSNE.Paused:
                 self.runbutton.setText("Resume")
             self.__timer.stop()
+            self.__timer_draw.stop()
 
     def __next_step(self):
         if self.__update_loop is None:
@@ -342,8 +357,10 @@ def __next_step(self):
         else:
             self.progressBarSet(100.0 * progress, processEvents=None)
             self.projection = projection
-            self.graph.update_coordinates()
-            self.graph.update_density()
+            if progress == 1 or self.needs_to_draw:
+                self.graph.update_coordinates()
+                self.graph.update_density()
+                self.needs_to_draw = False
             # schedule next update
             self.__timer.start()
 
@@ -360,15 +377,16 @@ def commit(self):
     def _get_projection_data(self):
         if self.data is None:
             return None
-        if self.projection is None:
-            variables = self._get_projection_variables()
-        else:
-            variables = self.projection.domain.attributes
         data = self.data.transform(
             Domain(self.data.domain.attributes,
                    self.data.domain.class_vars,
-                   self.data.domain.metas + variables))
+                   self.data.domain.metas + self._get_projection_variables()))
         data.metas[:, -2:] = self.get_embedding()
+        if self.projection is not None:
+            data.domain = Domain(
+                self.data.domain.attributes,
+                self.data.domain.class_vars,
+                self.data.domain.metas + self.projection.domain.attributes)
         return data
 
     def send_preprocessor(self):

Orange/widgets/unsupervised/tests/test_owtsne.py

@@ -170,6 +170,31 @@ def test_normalize_data(self):
             self.assertFalse(self.widget.controls.normalize.isEnabled())
             normalize.assert_not_called()
 
+    @patch("Orange.projection.manifold.TSNEModel.optimize")
+    def test_exaggeration_is_passed_through_properly(self, optimize):
+        def _check_exaggeration(call, exaggeration):
+            # Check the last call to `optimize`, so we catch one during the
+            # regular regime
+            name, args, kwargs = call.mock_calls[-1]
+            self.assertIn("exaggeration", kwargs)
+            self.assertEqual(kwargs["exaggeration"], exaggeration)
+
+        # Set value to 1
+        self.widget.controls.exaggeration.setValue(1)
+        self.send_signal(self.widget.Inputs.data, self.data)
+        self.commit_and_wait()
+        _check_exaggeration(optimize, 1)
+
+        # Reset and clear state
+        optimize.reset_mock()
+        self.send_signal(self.widget.Inputs.data, None)
+
+        # Change to 3
+        self.widget.controls.exaggeration.setValue(3)
+        self.send_signal(self.widget.Inputs.data, self.data)
+        self.commit_and_wait()
+        _check_exaggeration(optimize, 3)
+
 
 if __name__ == '__main__':
     unittest.main()