Save comparison multi plot image values across sessions (#4476)

extension/src/persistence/constants.ts

@@ -18,6 +18,7 @@ export enum PersistenceKey {
   PLOT_SECTION_COLLAPSED = 'plotSectionCollapsed:',
   PLOT_SELECTED_METRICS = 'plotSelectedMetrics:',
   PLOTS_SMOOTH_PLOT_VALUES = 'plotSmoothPlotValues:',
+  PLOTS_COMPARISON_MULTI_PLOT_VALUES = 'plotComparisonMultiPlotValues:',
   PLOT_TEMPLATE_ORDER = 'plotTemplateOrder:',
   SHOW_ONLY_CHANGED = 'columnsShowOnlyChanged:'
 }

extension/src/plots/model/index.ts

@@ -33,7 +33,8 @@ import {
   DEFAULT_NB_ITEMS_PER_ROW,
   PlotHeight,
   SmoothPlotValues,
-  ImagePlot
+  ImagePlot,
+  ComparisonMultiPlotValues
 } from '../webview/contract'
 import {
   EXPERIMENT_WORKSPACE_ID,
@@ -80,6 +81,7 @@ export class PlotsModel extends ModelWithPersistence {
 
   private comparisonData: ComparisonData = {}
   private comparisonOrder: string[]
+  private comparisonMultiPlotValues: ComparisonMultiPlotValues = {}
   private smoothPlotValues: SmoothPlotValues = {}
 
   private revisionData: RevisionData = {}
@@ -113,6 +115,10 @@ export class PlotsModel extends ModelWithPersistence {
       PersistenceKey.PLOTS_SMOOTH_PLOT_VALUES,
       {}
     )
+    this.comparisonMultiPlotValues = this.revive(
+      PersistenceKey.PLOTS_COMPARISON_MULTI_PLOT_VALUES,
+      {}
+    )
 
     this.cleanupOutdatedCustomPlotsState()
     this.cleanupOutdatedTrendsState()
@@ -307,6 +313,26 @@ export class PlotsModel extends ModelWithPersistence {
     this.persist(PersistenceKey.PLOT_COMPARISON_ORDER, this.comparisonOrder)
   }
 
+  public setComparisonMultiPlotValue(
+    revision: string,
+    path: string,
+    value: number
+  ) {
+    if (!this.comparisonMultiPlotValues[revision]) {
+      this.comparisonMultiPlotValues[revision] = {}
+    }
+
+    this.comparisonMultiPlotValues[revision][path] = value
+    this.persist(
+      PersistenceKey.PLOTS_COMPARISON_MULTI_PLOT_VALUES,
+      this.comparisonMultiPlotValues
+    )
+  }
+
+  public getComparisonMultiPlotValues() {
+    return this.comparisonMultiPlotValues
+  }
+
   public getSelectedRevisionIds() {
     return this.experiments.getSelectedRevisions().map(({ id }) => id)
   }
@@ -394,6 +420,7 @@ export class PlotsModel extends ModelWithPersistence {
       ...this.comparisonData,
       ...comparisonData
     }
+
     this.revisionData = {
       ...this.revisionData,
       ...revisionData

extension/src/plots/webview/contract.ts

@@ -66,11 +66,16 @@ export type Revision = {
   label: string
 }
 
+export type ComparisonMultiPlotValues = {
+  [revision: string]: { [path: string]: number }
+}
+
 export interface PlotsComparisonData {
   plots: ComparisonPlots
   width: number
   height: PlotHeight
   revisions: Revision[]
+  multiPlotValues: ComparisonMultiPlotValues
 }
 
 export type CustomPlotValues = {

extension/src/plots/webview/messages.ts

@@ -102,6 +102,12 @@ export class WebviewMessages {
         return this.selectPlotsFromWebview()
       case MessageFromWebviewType.SELECT_EXPERIMENTS:
         return this.selectExperimentsFromWebview()
+      case MessageFromWebviewType.SET_COMPARISON_MULTI_PLOT_VALUE:
+        return this.setComparisonMultiPlotValue(
+          message.payload.revision,
+          message.payload.path,
+          message.payload.value
+        )
       case MessageFromWebviewType.REMOVE_CUSTOM_PLOTS:
         return commands.executeCommand(
           RegisteredCommands.PLOTS_CUSTOM_REMOVE,
@@ -224,6 +230,20 @@ export class WebviewMessages {
     )
   }
 
+  private setComparisonMultiPlotValue(
+    revision: string,
+    path: string,
+    value: number
+  ) {
+    this.plots.setComparisonMultiPlotValue(revision, path, value)
+    this.sendComparisonPlots()
+    sendTelemetryEvent(
+      EventName.VIEWS_PLOTS_SET_COMPARISON_MULTI_PLOT_VALUE,
+      undefined,
+      undefined
+    )
+  }
+
   private setTemplateOrder(order: PlotsTemplatesReordered) {
     this.paths.setTemplateOrder(order)
     this.sendTemplatePlots()
@@ -345,6 +365,7 @@ export class WebviewMessages {
 
     return {
       height: this.plots.getHeight(PlotsSection.COMPARISON_TABLE),
+      multiPlotValues: this.plots.getComparisonMultiPlotValues(),
       plots: comparison.map(({ path, revisions }) => {
         return { path, revisions: this.getRevisionsWithCorrectUrls(revisions) }
       }),

extension/src/telemetry/constants.ts

@@ -92,6 +92,8 @@ export const EventName = Object.assign(
     VIEWS_PLOTS_SECTION_TOGGLE: 'views.plots.toggleSection',
     VIEWS_PLOTS_SELECT_EXPERIMENTS: 'view.plots.selectExperiments',
     VIEWS_PLOTS_SELECT_PLOTS: 'view.plots.selectPlots',
+    VIEWS_PLOTS_SET_COMPARISON_MULTI_PLOT_VALUE:
+      'view.plots.setComparisonMultiPlotValue',
     VIEWS_PLOTS_SET_SMOOTH_PLOT_VALUE: 'view.plots.setSmoothPlotValues',
     VIEWS_PLOTS_ZOOM_PLOT: 'views.plots.zoomPlot',
     VIEWS_REORDER_PLOTS_CUSTOM: 'views.plots.customReordered',
@@ -296,6 +298,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_PLOTS_ZOOM_PLOT]: { isImage: boolean }
   [EventName.VIEWS_REORDER_PLOTS_CUSTOM]: undefined
   [EventName.VIEWS_REORDER_PLOTS_TEMPLATES]: undefined
+  [EventName.VIEWS_PLOTS_SET_COMPARISON_MULTI_PLOT_VALUE]: undefined
   [EventName.VIEWS_PLOTS_SET_SMOOTH_PLOT_VALUE]: undefined
 
   [EventName.VIEWS_PLOTS_PATH_TREE_OPENED]: DvcRootCount

extension/src/test/fixtures/plotsDiff/index.ts

@@ -811,6 +811,7 @@ export const getComparisonWebviewMessage = (
 
   return {
     revisions: getRevisions(),
+    multiPlotValues: {},
     plots: Object.values(plotAcc),
     width: DEFAULT_PLOT_WIDTH,
     height: DEFAULT_PLOT_HEIGHT

extension/src/test/suite/plots/index.test.ts

@@ -1065,6 +1065,42 @@ suite('Plots Test Suite', () => {
       )
     })
 
+    it('should handle an update comparison multi plot value message from the webview', async () => {
+      const { mockMessageReceived, plotsModel } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+      const multiImg = comparisonPlotsFixture.plots[3]
+
+      const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockSetComparisonMultiPlotValue = stub(
+        plotsModel,
+        'setComparisonMultiPlotValue'
+      )
+
+      mockMessageReceived.fire({
+        payload: {
+          path: multiImg.path,
+          revision: 'main',
+          value: 5
+        },
+        type: MessageFromWebviewType.SET_COMPARISON_MULTI_PLOT_VALUE
+      })
+
+      expect(mockSendTelemetryEvent).to.be.called
+      expect(mockSendTelemetryEvent).to.be.calledWithExactly(
+        EventName.VIEWS_PLOTS_SET_COMPARISON_MULTI_PLOT_VALUE,
+        undefined,
+        undefined
+      )
+      expect(mockSetComparisonMultiPlotValue).to.be.called
+      expect(mockSetComparisonMultiPlotValue).to.be.calledWithExactly(
+        'main',
+        multiImg.path,
+        5
+      )
+    })
+
     it('should handle the CLI throwing an error', async () => {
       const { data, errorsModel, mockPlotsDiff, plots, plotsModel } =
         await buildPlots({ disposer: disposable, plotsDiff: plotsDiffFixture })

extension/src/webview/contract.ts

@@ -44,6 +44,7 @@ export enum MessageFromWebviewType {
   RESIZE_COLUMN = 'resize-column',
   RESIZE_PLOTS = 'resize-plots',
   SAVE_STUDIO_TOKEN = 'save-studio-token',
+  SET_COMPARISON_MULTI_PLOT_VALUE = 'update-comparison-multi-plot-value',
   SET_SMOOTH_PLOT_VALUE = 'update-smooth-plot-value',
   SHOW_EXPERIMENT_LOGS = 'show-experiment-logs',
   SHOW_WALKTHROUGH = 'show-walkthrough',
@@ -211,6 +212,10 @@ export type MessageFromWebview =
       type: MessageFromWebviewType.REORDER_PLOTS_COMPARISON_ROWS
       payload: string[]
     }
+  | {
+      type: MessageFromWebviewType.SET_COMPARISON_MULTI_PLOT_VALUE
+      payload: { path: string; revision: string; value: number }
+    }
   | {
       type: MessageFromWebviewType.REORDER_PLOTS_CUSTOM
       payload: string[]

webview/src/plots/components/App.test.tsx

@@ -225,6 +225,7 @@ describe('App', () => {
     renderAppWithOptionalData({
       comparison: {
         height: DEFAULT_PLOT_HEIGHT,
+        multiPlotValues: {},
         plots: [
           {
             path: 'training/plots/images/misclassified.jpg',
@@ -279,6 +280,7 @@ describe('App', () => {
     renderAppWithOptionalData({
       comparison: {
         height: DEFAULT_PLOT_HEIGHT,
+        multiPlotValues: {},
         plots: [
           {
             path: 'training/plots/images/image',
@@ -1757,6 +1759,25 @@ describe('App', () => {
 
       const workspaceImgs =
         comparisonTableFixture.plots[3].revisions.workspace.imgs
+      const multiImgPlots = screen.getAllByTestId('multi-image-cell')
+      const slider = within(multiImgPlots[0]).getByRole('slider')
+      const workspaceImgEl = within(multiImgPlots[0]).getByRole('img')
+
+      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[0].url)
+
+      fireEvent.change(slider, { target: { value: 3 } })
+
+      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[3].url)
+    })
+
+    it('should send a message when the slider changes', async () => {
+      renderAppWithOptionalData({
+        comparison: comparisonTableFixture
+      })
+
+      const multiImg = comparisonTableFixture.plots[3]
+      const workspacePlot = multiImg.revisions.workspace
+      const workspaceImgs = workspacePlot.imgs
 
       const multiImgPlots = screen.getAllByTestId('multi-image-cell')
       const slider = within(multiImgPlots[0]).getByRole('slider')
@@ -1766,6 +1787,34 @@ describe('App', () => {
 
       fireEvent.change(slider, { target: { value: 3 } })
 
+      await waitFor(
+        () => {
+          expect(mockPostMessage).toHaveBeenCalledWith({
+            payload: {
+              path: multiImg.path,
+              revision: workspacePlot.id,
+              value: 3
+            },
+            type: MessageFromWebviewType.SET_COMPARISON_MULTI_PLOT_VALUE
+          })
+        },
+        { timeout: 1000 }
+      )
+    })
+
+    it('should set default slider value if given a saved value', () => {
+      const multiImg = comparisonTableFixture.plots[3]
+      renderAppWithOptionalData({
+        comparison: {
+          ...comparisonTableFixture,
+          multiPlotValues: { workspace: { [multiImg.path]: 3 } }
+        }
+      })
+
+      const workspaceImgs = multiImg.revisions.workspace.imgs
+      const multiImgPlots = screen.getAllByTestId('multi-image-cell')
+      const workspaceImgEl = within(multiImgPlots[0]).getByRole('img')
+
       expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[3].url)
     })
 
@@ -1795,7 +1844,6 @@ describe('App', () => {
       })
 
       const mainImgs = comparisonTableFixture.plots[3].revisions.main.imgs
-
       const multiImgPlots = screen.getAllByTestId('multi-image-cell')
       const slider = within(multiImgPlots[1]).getByRole('slider')
 

webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx

@@ -1,16 +1,23 @@
-import React, { useCallback, useState } from 'react'
-import { useDispatch } from 'react-redux'
+import React, { useEffect, useCallback, useRef, useState } from 'react'
+import { useDispatch, useSelector } from 'react-redux'
 import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
 import { ComparisonTableCell } from './ComparisonTableCell'
 import styles from '../styles.module.scss'
 import { changeDisabledDragIds } from '../comparisonTableSlice'
+import { setComparisonMultiPlotValue } from '../../../util/messages'
+import { PlotsState } from '../../../store'
 
 export const ComparisonTableMultiCell: React.FC<{
   path: string
   plot: ComparisonPlot
 }> = ({ path, plot }) => {
-  const [currentStep, setCurrentStep] = useState<number>(0)
+  const values = useSelector(
+    (state: PlotsState) => state.comparison.multiPlotValues
+  )
+  const [currentStep, setCurrentStep] = useState(values?.[plot.id]?.[path] || 0)
   const dispatch = useDispatch()
+  const maxStep = plot.imgs.length - 1
+  const changeDebounceTimer = useRef(0)
 
   const addDisabled = useCallback(() => {
     dispatch(changeDisabledDragIds([path]))
@@ -20,6 +27,16 @@ export const ComparisonTableMultiCell: React.FC<{
     dispatch(changeDisabledDragIds([]))
   }, [dispatch])
 
+  useEffect(() => {
+    window.clearTimeout(changeDebounceTimer.current)
+    changeDebounceTimer.current = window.setTimeout(() => {
+      if (currentStep === values?.[plot.id]?.[path]) {
+        return
+      }
+      setComparisonMultiPlotValue(path, plot.id, currentStep)
+    }, 500)
+  }, [values, path, plot.id, currentStep])
+
   return (
     <div data-testid="multi-image-cell" className={styles.multiImageWrapper}>
       <ComparisonTableCell
@@ -45,10 +62,14 @@ export const ComparisonTableMultiCell: React.FC<{
         <input
           name={`${plot.id}-step`}
           min="0"
-          max={plot.imgs.length - 1}
-          value={currentStep}
+          max={maxStep}
           type="range"
+          value={currentStep}
           onChange={event => {
+            if (!event.target) {
+              return
+            }
+
             setCurrentStep(Number(event.target.value))
           }}
         />

webview/src/plots/components/comparisonTable/comparisonTableSlice.ts

@@ -21,6 +21,7 @@ export const comparisonTableInitialState: ComparisonTableState = {
   hasData: false,
   height: DEFAULT_HEIGHT[PlotsSection.COMPARISON_TABLE],
   isCollapsed: DEFAULT_SECTION_COLLAPSED[PlotsSection.COMPARISON_TABLE],
+  multiPlotValues: {},
   plots: [],
   revisions: [],
   rowHeight: DEFAULT_ROW_HEIGHT,

webview/src/plots/util/messages.ts

@@ -96,6 +96,17 @@ export const selectRevisions = () =>
     type: MessageFromWebviewType.SELECT_EXPERIMENTS
   })
 
+export const setComparisonMultiPlotValue = (
+  path: string,
+  revision: string,
+  value: number
+) => {
+  sendMessage({
+    payload: { path, revision, value },
+    type: MessageFromWebviewType.SET_COMPARISON_MULTI_PLOT_VALUE
+  })
+}
+
 export const togglePlotsSection = (
   sectionKey: PlotsSection,
   sectionCollapsed: boolean

webview/src/stories/ComparisonTable.stories.tsx

@@ -47,6 +47,7 @@ const Template: StoryFn = ({ plots, revisions }) => {
       <MockedState
         data={{
           height: DEFAULT_PLOT_HEIGHT,
+          multiPlotValues: {},
           plots,
           revisions,
           width: DEFAULT_NB_ITEMS_PER_ROW