Embedding Projector: fix bookmark projection state

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. */

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

@@ -190,7 +190,7 @@ class BookmarkPanel extends LegacyElementMixin(PolymerElement) {
     this.notifyPath(path, selected);
   }
   /**
-   * Return an event's to their bookmark index.
+   * Returns the bookmark index of the event.
    */
   private getBookmarkIndex(evt: any) {
     return evt.model.__data.index;

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);