From 170100dc0a27ca4fb9f63f86cab8adf31968542f Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 31 May 2022 16:25:32 +1000
Subject: [PATCH] add legend for zoomed in checkpoint plots

---
 webview/src/plots/components/ZoomablePlot.tsx |  25 +--
 .../plots/components/checkpointPlots/util.ts  | 150 +++++++++---------
 2 files changed, 93 insertions(+), 82 deletions(-)

diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index 58e3747a77..d8700d63cb 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -16,6 +16,20 @@ interface ZoomablePlotOwnProps extends ZoomablePlotProps {
   id: string
 }
 
+const removeLegendSuppression = (plotProps: VegaLiteProps): VegaLiteProps => {
+  const plot = { ...plotProps } as VegaLiteProps & {
+    spec: {
+      encoding?: {
+        color?: { legend?: { disable?: boolean } }
+      }
+    }
+  }
+  if (plot.spec.encoding?.color?.legend?.disable !== undefined) {
+    delete plot.spec.encoding.color.legend.disable
+  }
+  return plot
+}
+
 export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   plotProps,
   renderZoomedInPlot,
@@ -32,15 +46,8 @@ export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   }, [plotProps, id, renderZoomedInPlot])
 
   const handleOnClick = () => {
-    const zoomedPlotWithLegend = { ...plotProps } as VegaLiteProps & {
-      spec: { encoding?: { color?: { legend?: { disable?: boolean } } } }
-    }
-    if (
-      zoomedPlotWithLegend.spec.encoding?.color?.legend?.disable !== undefined
-    ) {
-      delete zoomedPlotWithLegend.spec.encoding.color.legend.disable
-    }
-    return renderZoomedInPlot(zoomedPlotWithLegend, id)
+    const zoomedPlot = removeLegendSuppression(plotProps)
+    return renderZoomedInPlot(zoomedPlot, id)
   }
 
   return (
diff --git a/webview/src/plots/components/checkpointPlots/util.ts b/webview/src/plots/components/checkpointPlots/util.ts
index 2c6ad57def..94e099555b 100644
--- a/webview/src/plots/components/checkpointPlots/util.ts
+++ b/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)