Embedding Projector: fix bookmark projection state

alicialics · alicialics · commit 7bf028f39c40 · 2023-04-15T18:19:29.000Z
diff --git a/tensorboard/plugins/projector/vz_projector/data.ts b/tensorboard/plugins/projector/vz_projector/data.ts
@@ -680,6 +680,8 @@ export class State {
   filteredPoints: number[];
   /** The indices of selected points. */
   selectedPoints: number[] = [];
+  /** The shuffled indices of points. */
+  shuffledDataIndices: number[] = [];
   /** Camera state (2d/3d, position, target, zoom, etc). */
   cameraDef: CameraDef;
   /** Color by option. */
diff --git a/tensorboard/plugins/projector/vz_projector/vz-projector.ts b/tensorboard/plugins/projector/vz_projector/vz-projector.ts
@@ -650,6 +650,7 @@ class Projector
     state.tSNEIteration = this.dataSet.tSNEIteration;
     state.selectedPoints = this.selectedPointIndices;
     state.filteredPoints = this.dataSetFilterIndices!;
+    state.shuffledDataIndices = this.dataSet.shuffledDataIndices;
     this.projectorScatterPlotAdapter.populateBookmarkFromUI(state);
     state.selectedColorOptionName = this.dataPanel.selectedColorOptionName;
     state.forceCategoricalColoring = this.dataPanel.forceCategoricalColoring;
@@ -674,10 +675,14 @@ class Projector
       const point = this.dataSet.points[i];
       const projection = state.projections[i];
       const keys = Object.keys(projection);
+      point.projections = {};
       for (let j = 0; j < keys.length; ++j) {
         point.projections[keys[j]] = projection[keys[j]];
       }
     }
+    if (state.shuffledDataIndices) {
+      this.dataSet.shuffledDataIndices = state.shuffledDataIndices;
+    }
     this.dataSet.hasTSNERun = state.selectedProjection === 'tsne';
     this.dataSet.tSNEIteration = state.tSNEIteration;
     this.projectionsPanel.restoreUIFromBookmark(state);
