From 6436b3b167cd455e4e82920cf6083703d4fc6e74 Mon Sep 17 00:00:00 2001
From: Zoey Li <31378877+alicialics@users.noreply.github.com>
Date: Sat, 15 Apr 2023 18:16:03 +0000
Subject: [PATCH] Embedding Projector: fix bookmark projection state

---
 tensorboard/plugins/projector/vz_projector/data.ts           | 2 ++
 .../projector/vz_projector/vz-projector-bookmark-panel.ts    | 2 +-
 tensorboard/plugins/projector/vz_projector/vz-projector.ts   | 5 +++++
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/tensorboard/plugins/projector/vz_projector/data.ts b/tensorboard/plugins/projector/vz_projector/data.ts
index 9cb434ef2f9..efd9af6da32 100644
--- 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-bookmark-panel.ts b/tensorboard/plugins/projector/vz_projector/vz-projector-bookmark-panel.ts
index a46290d9132..43c60760adb 100644
--- a/tensorboard/plugins/projector/vz_projector/vz-projector-bookmark-panel.ts
+++ b/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;
diff --git a/tensorboard/plugins/projector/vz_projector/vz-projector.ts b/tensorboard/plugins/projector/vz_projector/vz-projector.ts
index 49150984071..f17de05c750 100644
--- 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);