Embedding Projector: fix tSNE tweaking hyperparams

Author: alicialics

Diff for: tensorboard/plugins/projector/vz_projector/bh_tsne.ts

@@ -299,10 +299,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() {

Diff for: tensorboard/plugins/projector/vz_projector/data.ts

@@ -334,6 +334,19 @@ export class DataSet {
     this.tSNEIteration = 0;
     let sampledIndices = this.shuffledDataIndices.slice(0, TSNE_SAMPLE_SIZE);
     let step = () => {
+      if (!this.tsne.getSolution()) {
+        const rng = this.tsne.getRng();
+        const d = this.tsne.getDim();
+        const solution = new Float64Array(sampledIndices.length * d);
+        sampledIndices.forEach((index, i) => {
+          let dataPoint = this.points[index];
+          for (let j = 0; j < d; j++) {
+            solution[i * d + j] = dataPoint.projections[`tsne-${j}`] ?? rng();
+          }
+        });
+        this.tsne.setSolution(solution);
+        return;
+      }
       if (this.tSNEShouldStop) {
         this.projections['tsne'] = false;
         stepCallback(null!);
@@ -344,12 +357,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;

Diff for: 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;
@@ -380,6 +384,9 @@ class ProjectionsPanel extends LegacyElementMixin(PolymerElement) {
   }
   @observe('tSNEis3d')
   _tsneDimensionToggleObserver() {
+    if (this.dataSet?.hasTSNERun) {
+      this.dataSet.hasTSNERun = false;
+    }
     this.beginProjection(this.currentProjection);
   }
   @observe('umapIs3d')