Add legend to zoomed in plots

webview/package.json

@@ -25,6 +25,8 @@
     "@tippyjs/react": "^4.2.6",
     "@vscode/webview-ui-toolkit": "^1.0.0",
     "classnames": "^2.2.6",
+    "lodash.clonedeep": "^4.5.0",
+    "lodash.merge": "^4.6.2",
     "react": "^17.0.1",
     "react-dom": "^17.0.1",
     "react-table": "^7.7.0",

webview/src/plots/components/Plots.tsx

@@ -1,11 +1,10 @@
 import { PlotSize, Section } from 'dvc/src/plots/webview/contract'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 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 { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import { PlotsSizeProvider } from './PlotsSizeContext'
 import { AddPlots, Welcome } from './GetStarted'
+import { ZoomedInPlot } from './ZoomedInPlot'
 import { CheckpointPlotsWrapper } from './checkpointPlots/CheckpointPlotsWrapper'
 import { TemplatePlotsWrapper } from './templatePlots/TemplatePlotsWrapper'
 import { ComparisonTableWrapper } from './comparisonTable/ComparisonTableWrapper'
@@ -16,7 +15,6 @@ import { Modal } from '../../shared/components/modal/Modal'
 import { WebviewWrapper } from '../../shared/components/webviewWrapper/WebviewWrapper'
 import { DragDropProvider } from '../../shared/components/dragDrop/DragDropContext'
 import { sendMessage } from '../../shared/vscode'
-import { getThemeValue, ThemeProperty } from '../../util/styles'
 import { GetStarted } from '../../shared/components/getStarted/GetStarted'
 
 interface PlotsProps {
@@ -153,21 +151,7 @@ const PlotsContent = ({ state }: PlotsProps) => {
 
       {zoomedInPlot && (
         <Modal onClose={handleModalClose}>
-          <div className={styles.zoomedInPlot} data-testid="zoomed-in-plot">
-            <VegaLite
-              {...zoomedInPlot}
-              config={{
-                ...(zoomedInPlot.config as Config),
-                background: getThemeValue(ThemeProperty.MENU_BACKGROUND)
-              }}
-              actions={{
-                compiled: false,
-                editor: false,
-                export: true,
-                source: false
-              }}
-            />
-          </div>
+          <ZoomedInPlot props={zoomedInPlot} />
         </Modal>
       )}
     </>

webview/src/plots/components/ZoomedInPlot.tsx

@@ -0,0 +1,36 @@
+import React from 'react'
+import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
+import { Config } from 'vega-lite'
+import merge from 'lodash.merge'
+import cloneDeep from 'lodash.clonedeep'
+import styles from './styles.module.scss'
+import { getThemeValue, ThemeProperty } from '../../util/styles'
+
+export type ZoomedInPlotProps = {
+  props: VegaLiteProps
+}
+
+export const ZoomedInPlot: React.FC<ZoomedInPlotProps> = ({
+  props
+}: ZoomedInPlotProps) => (
+  <div className={styles.zoomedInPlot} data-testid="zoomed-in-plot">
+    <VegaLite
+      {...merge(
+        { ...cloneDeep(props) },
+        {
+          spec: { encoding: { color: { legend: { disable: false } } } }
+        }
+      )}
+      config={{
+        ...(props.config as Config),
+        background: getThemeValue(ThemeProperty.MENU_BACKGROUND)
+      }}
+      actions={{
+        compiled: false,
+        editor: false,
+        export: true,
+        source: false
+      }}
+    />
+  </div>
+)

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/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)
+}