Embedding Projector: fix tSNE tweaking hyperparams

Author: alicialics

tensorboard/plugins/projector/vz_projector/bh_tsne.ts

@@ -289,8 +289,10 @@ export class TSNE {
   }
   // (re)initializes the solution to random
   initSolution() {
-    // generate random solution to t-SNE
-    this.Y = randnMatrix(this.N, this.dim, this.rng); // the solution
+    if (!this.Y) {
+      // generate random solution to t-SNE
+      this.Y = randnMatrix(this.N, this.dim, this.rng); // the solution
+    }
     this.gains = arrayofs(this.N, this.dim, 1); // step gains
     // to accelerate progress in unchanging directions
     this.ystep = arrayofs(this.N, this.dim, 0); // momentum accumulator
@@ -299,10 +301,16 @@ export class TSNE {
   getDim() {
     return this.dim;
   }
+  getRng() {
+    return this.rng;
+  }
   // return pointer to current solution
   getSolution() {
     return this.Y;
   }
+  setSolution(solution: Float64Array) {
+    this.Y = solution;
+  }
   // For each point, randomly offset point within a 5% hypersphere centered
   // around it, whilst remaining in the assumed t-SNE plot hypersphere
   perturb() {

tensorboard/plugins/projector/vz_projector/data.ts

@@ -334,6 +334,10 @@ export class DataSet {
     this.tSNEIteration = 0;
     let sampledIndices = this.shuffledDataIndices.slice(0, TSNE_SAMPLE_SIZE);
     let step = () => {
+      if (!this.tsne.getSolution()) {
+        this.initTsneSolutionFromCurrentProjection();
+        return;
+      }
       if (this.tSNEShouldStop) {
         this.projections['tsne'] = false;
         stepCallback(null!);
@@ -344,12 +348,11 @@ export class DataSet {
       if (!this.tSNEShouldPause) {
         this.tsne.step();
         let result = this.tsne.getSolution();
+        const d = this.tsne.getDim();
         sampledIndices.forEach((index, i) => {
           let dataPoint = this.points[index];
-          dataPoint.projections['tsne-0'] = result[i * tsneDim + 0];
-          dataPoint.projections['tsne-1'] = result[i * tsneDim + 1];
-          if (tsneDim === 3) {
-            dataPoint.projections['tsne-2'] = result[i * tsneDim + 2];
+          for (let j = 0; j < d; j++) {
+            dataPoint.projections[`tsne-${j}`] = result[i * d + j];
           }
         });
         this.projections['tsne'] = true;
@@ -486,6 +489,20 @@ export class DataSet {
       return Promise.resolve(result);
     }
   }
+  /* initialize the new tsne solution from current projection data */
+  initTsneSolutionFromCurrentProjection() {
+    const sampledIndices = this.shuffledDataIndices.slice(0, TSNE_SAMPLE_SIZE);
+    const rng = this.tsne.getRng();
+    const d = this.tsne.getDim();
+    const solution = new Float64Array(sampledIndices.length * d);
+    sampledIndices.forEach((index, i) => {
+      const dataPoint = this.points[index];
+      for (let j = 0; j < d; j++) {
+        solution[i * d + j] = dataPoint.projections[`tsne-${j}`] ?? rng();
+      }
+    });
+    this.tsne.setSolution(solution);
+  }
   /* Perturb TSNE and update dataset point coordinates. */
   perturbTsne() {
     if (this.hasTSNERun && this.tsne) {

tensorboard/plugins/projector/vz_projector/vz-projector-projections-panel.ts

@@ -154,6 +154,10 @@ class ProjectionsPanel extends LegacyElementMixin(PolymerElement) {
   private updateTSNEPerplexityFromSliderChange() {
     if (this.perplexitySlider) {
       this.perplexity = +this.perplexitySlider.value;
+      if (this.dataSet?.hasTSNERun) {
+        this.dataSet.hasTSNERun = false;
+        this.beginProjection(this.currentProjection);
+      }
     }
     (this.$$('.tsne-perplexity span') as HTMLSpanElement).innerText =
       '' + this.perplexity;
@@ -295,6 +299,10 @@ class ProjectionsPanel extends LegacyElementMixin(PolymerElement) {
     }
     if (bookmark.selectedProjection != null) {
       this.showTab(bookmark.selectedProjection);
+      if (this.currentProjection === 'tsne') {
+        this.runTSNE();
+        this.dataSet.initTsneSolutionFromCurrentProjection();
+      }
     }
     this.enablePolymerChangesTriggerReprojection();
   }
@@ -380,6 +388,9 @@ class ProjectionsPanel extends LegacyElementMixin(PolymerElement) {
   }
   @observe('tSNEis3d')
   _tsneDimensionToggleObserver() {
+    if (this.dataSet?.hasTSNERun) {
+      this.dataSet.hasTSNERun = false;
+    }
     this.beginProjection(this.currentProjection);
   }
   @observe('umapIs3d')