Add legend to zoomed in plots

webview/src/plots/components/Plots.tsx

@@ -4,6 +4,7 @@ import React, { useEffect, useRef, useState, useCallback } from 'react'
 import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import { Config } from 'vega-lite'
 import styles from './styles.module.scss'
+import { removeLegendSuppression } from './util'
 import { PlotsSizeProvider } from './PlotsSizeContext'
 import { AddPlots, Welcome } from './GetStarted'
 import { CheckpointPlotsWrapper } from './checkpointPlots/CheckpointPlotsWrapper'
@@ -43,13 +44,15 @@ const PlotsContent = ({ state }: PlotsProps) => {
   const handleZoomInPlot = useCallback(
     (props: VegaLiteProps, id: string, refresh?: boolean) => {
       if (!refresh) {
-        setZoomedInPlot(props)
+        setZoomedInPlot(removeLegendSuppression(props))
         zoomedInPlotId.current = id
         return
       }
 
       if (zoomedInPlotId.current === id) {
-        setZoomedInPlot(plot => (plot ? props : undefined))
+        setZoomedInPlot(plot =>
+          plot ? removeLegendSuppression(props) : undefined
+        )
       }
     },
     [setZoomedInPlot]

webview/src/plots/components/ZoomablePlot.tsx

@@ -1,6 +1,7 @@
 import React, { useEffect, useRef } from 'react'
 import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import styles from './styles.module.scss'
+import { removeLegendSuppression } from './util'
 import { GripIcon } from '../../shared/components/dragDrop/GripIcon'
 
 export interface ZoomablePlotProps {
@@ -16,16 +17,21 @@ interface ZoomablePlotOwnProps extends ZoomablePlotProps {
   id: string
 }
 
+const areDifferentPlotProps = (
+  currentPlotProps: VegaLiteProps,
+  plotProps: VegaLiteProps
+) =>
+  JSON.stringify(removeLegendSuppression(currentPlotProps)) !==
+  JSON.stringify(removeLegendSuppression(plotProps))
+
 export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   plotProps,
   renderZoomedInPlot,
   id
 }) => {
   const previousPlotProps = useRef(plotProps)
   useEffect(() => {
-    if (
-      JSON.stringify(previousPlotProps.current) !== JSON.stringify(plotProps)
-    ) {
+    if (areDifferentPlotProps(previousPlotProps.current, plotProps)) {
       renderZoomedInPlot(plotProps, id, true)
       previousPlotProps.current = plotProps
     }

webview/src/plots/components/checkpointPlots/util.ts

@@ -4,86 +4,90 @@ import { ColorScale } from 'dvc/src/plots/webview/contract'
 export const createSpec = (
   title: string,
   scale?: ColorScale
-): VisualizationSpec => ({
-  $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
-  data: { name: 'values' },
-  encoding: {
-    x: {
-      axis: { format: '0d', tickMinStep: 1 },
-      field: 'iteration',
-      title: 'iteration',
-      type: 'quantitative'
+): VisualizationSpec =>
+  ({
+    $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
+    data: { name: 'values' },
+    encoding: {
+      color: {
+        field: 'group',
+        legend: { disable: true },
+        scale,
+        title: 'rev',
+        type: 'nominal'
+      },
+      x: {
+        axis: { format: '0d', tickMinStep: 1 },
+        field: 'iteration',
+        title: 'iteration',
+        type: 'quantitative'
+      },
+      y: {
+        field: 'y',
+        scale: { zero: false },
+        title,
+        type: 'quantitative'
+      }
     },
-    y: {
-      field: 'y',
-      scale: { zero: false },
-      title,
-      type: 'quantitative'
-    }
-  },
-  height: 'container',
-  layer: [
-    {
-      encoding: {
-        color: { field: 'group', legend: null, scale, type: 'nominal' }
+    height: 'container',
+    layer: [
+      {
+        layer: [
+          { mark: { type: 'line' } },
+          {
+            mark: { type: 'point' },
+            transform: [
+              {
+                filter: { empty: false, param: 'hover' }
+              }
+            ]
+          }
+        ]
       },
-
-      layer: [
-        { mark: { type: 'line' } },
-        {
-          mark: { type: 'point' },
-          transform: [
+      {
+        encoding: {
+          opacity: { value: 0 },
+          tooltip: [
+            { field: 'group', title: 'name' },
             {
-              filter: { empty: false, param: 'hover' }
+              field: 'y',
+              title: title.slice(Math.max(0, title.indexOf(':') + 1)),
+              type: 'quantitative'
             }
           ]
-        }
-      ]
-    },
-    {
-      encoding: {
-        opacity: { value: 0 },
-        tooltip: [
-          { field: 'group', title: 'name' },
+        },
+        mark: { type: 'rule' },
+        params: [
           {
-            field: 'y',
-            title: title.slice(Math.max(0, title.indexOf(':') + 1)),
-            type: 'quantitative'
+            name: 'hover',
+            select: {
+              clear: 'mouseout',
+              fields: ['iteration', 'y'],
+              nearest: true,
+              on: 'mouseover',
+              type: 'point'
+            }
           }
         ]
       },
-      mark: { type: 'rule' },
-      params: [
-        {
-          name: 'hover',
-          select: {
-            clear: 'mouseout',
-            fields: ['iteration', 'y'],
-            nearest: true,
-            on: 'mouseover',
-            type: 'point'
+      {
+        encoding: {
+          color: { field: 'group', scale },
+          x: { aggregate: 'max', field: 'iteration', type: 'quantitative' },
+          y: {
+            aggregate: { argmax: 'iteration' },
+            field: 'y',
+            type: 'quantitative'
           }
-        }
-      ]
-    },
-    {
-      encoding: {
-        color: { field: 'group', scale },
-        x: { aggregate: 'max', field: 'iteration', type: 'quantitative' },
-        y: {
-          aggregate: { argmax: 'iteration' },
-          field: 'y',
-          type: 'quantitative'
-        }
-      },
-      mark: { stroke: null, type: 'circle' }
-    }
-  ],
-  transform: [
-    {
-      as: 'y',
-      calculate: "format(datum['y'],'.5f')"
-    }
-  ],
-  width: 'container'
-})
+        },
+        mark: { stroke: null, type: 'circle' }
+      }
+    ],
+    transform: [
+      {
+        as: 'y',
+        calculate: "format(datum['y'],'.5f')"
+      }
+    ],
+    width: 'container'
+  } as VisualizationSpec)

webview/src/plots/components/util.ts

@@ -1,4 +1,5 @@
 import { PlotSize } from 'dvc/src/plots/webview/contract'
+import { VegaLiteProps } from 'react-vega/lib/VegaLite'
 
 const MaxItemsBeforeVirtualization = {
   [PlotSize.LARGE]: 10,
@@ -44,3 +45,19 @@ export const getNbItemsPerRow = (size: PlotSize) => {
 
 export const shouldUseVirtualizedGrid = (nbItems: number, size: PlotSize) =>
   nbItems > MaxItemsBeforeVirtualization[size]
+
+export const removeLegendSuppression = (
+  plotProps?: VegaLiteProps
+): VegaLiteProps => {
+  const plot = { ...plotProps } as VegaLiteProps & {
+    spec: {
+      encoding?: {
+        color?: { legend?: { disable?: boolean } }
+      }
+    }
+  }
+  if (plot?.spec?.encoding?.color?.legend?.disable !== undefined) {
+    plot.spec.encoding.color.legend.disable = false
+  }
+  return plot
+}

webview/src/stories/Plots.stories.tsx

@@ -185,3 +185,16 @@ MultiviewZoomedInPlot.play = async ({ canvasElement }) => {
 
   fireEvent.click(plotButton)
 }
+
+export const CheckpointZoomedInPlot = Template.bind({})
+CheckpointZoomedInPlot.parameters = {
+  chromatic: { delay: 500 }
+}
+CheckpointZoomedInPlot.play = async ({ canvasElement }) => {
+  const canvas = within(canvasElement)
+  const plot = await canvas.findByText('summary.json:val_accuracy')
+
+  plot.scrollIntoView()
+
+  fireEvent.click(plot)
+}