Ensure custom plots legend matches plot values (#4729)

Author: julieg18

extension/src/plots/model/collect.test.ts

@@ -66,6 +66,34 @@ describe('collectCustomPlots', () => {
     })
     expect(data[0].values.slice(-1)[0].id).toStrictEqual('main')
   })
+
+  it('should create custom plot scales that match the collected values', () => {
+    const expectedOutput: CustomPlotData[] = customPlotsFixture.plots
+    const data = collectCustomPlots({
+      colorScale: {
+        domain: ['main', 'exp-e7a67', 'test-branch', 'exp-83425', 'failed-exp'],
+        range: ['#13adc7', '#f46837', '#48bb78', '#4299e1', '#f56565']
+      },
+      experiments: [
+        ...experimentsWithCommits,
+        {
+          branch: 'main',
+          id: 'weird-exp',
+          label: 'exp with no metrics or params'
+        },
+        {
+          branch: 'main',
+          error: 'failed to run',
+          id: 'failed-exp',
+          label: '123'
+        }
+      ],
+      height: DEFAULT_PLOT_HEIGHT,
+      nbItemsPerRow: DEFAULT_NB_ITEMS_PER_ROW,
+      plotsOrderValues: customPlotsOrderFixture
+    })
+    expect(data).toStrictEqual(expectedOutput)
+  })
 })
 
 describe('collectData', () => {

extension/src/plots/model/collect.ts

@@ -80,19 +80,63 @@ const getValues = (
   return values
 }
 
+const removeSelectedExperiment = (
+  orderedColorScale: ColorScale,
+  hasValues: boolean,
+  idx: number
+) => {
+  const isSelectedExperiment = idx !== -1
+  if (!isSelectedExperiment || hasValues) {
+    return
+  }
+
+  orderedColorScale.domain.splice(idx, 1)
+  orderedColorScale.range.splice(idx, 1)
+}
+
+const fillColorScale = (
+  experiments: Experiment[],
+  colorScale: ColorScale | undefined,
+  valueIds: Set<string>
+) => {
+  const orderedColorScale = {
+    domain: [...(colorScale?.domain || [])],
+    range: [...(colorScale?.range || [])]
+  }
+
+  for (const experiment of experiments) {
+    const { id } = experiment
+    const idx = orderedColorScale.domain.indexOf(id)
+    const isSelectedExperiment = idx !== -1
+    const hasValues = valueIds.has(id)
+
+    if (!hasValues || isSelectedExperiment) {
+      removeSelectedExperiment(orderedColorScale, hasValues, idx)
+      continue
+    }
+
+    orderedColorScale.domain.push(id)
+    orderedColorScale.range.push('#4c78a8' as Color)
+  }
+
+  return orderedColorScale
+}
+
 const getCustomPlotData = (
   orderValue: CustomPlotsOrderValue,
   experiments: Experiment[],
   height: number,
   nbItemsPerRow: number,
-  completeColorScale: ColorScale,
-  renderLastIds: Set<string>
+  colorScale: ColorScale | undefined
 ): CustomPlotData => {
   const { metric, param } = orderValue
   const metricPath = getFullValuePath(ColumnType.METRICS, metric)
   const paramPath = getFullValuePath(ColumnType.PARAMS, param)
 
+  const renderLastIds = new Set(colorScale?.domain)
   const values = getValues(experiments, metricPath, paramPath, renderLastIds)
+  const valueIds = new Set(values.map(({ id }) => id))
+  const completeColorScale = fillColorScale(experiments, colorScale, valueIds)
 
   const [{ param: paramVal, metric: metricVal }] = values
   const yTitle = truncateVerticalTitle(metric, nbItemsPerRow, height) as string
@@ -115,26 +159,6 @@ const getCustomPlotData = (
   } as CustomPlotData
 }
 
-const fillColorScale = (
-  colorScale: ColorScale | undefined,
-  experiments: Experiment[]
-) => {
-  const completeColorScale = {
-    domain: [...(colorScale?.domain || [])],
-    range: [...(colorScale?.range || [])]
-  }
-
-  for (const experiment of experiments) {
-    const { id } = experiment
-    if (completeColorScale.domain.includes(id)) {
-      continue
-    }
-    completeColorScale.domain.push(id)
-    completeColorScale.range.push('#4c78a8' as Color)
-  }
-  return completeColorScale
-}
-
 export const collectCustomPlots = ({
   colorScale,
   plotsOrderValues,
@@ -150,18 +174,9 @@ export const collectCustomPlots = ({
 }): CustomPlotData[] => {
   const plots = []
 
-  const completeColorScale = fillColorScale(colorScale, experiments)
-
   for (const value of plotsOrderValues) {
     plots.push(
-      getCustomPlotData(
-        value,
-        experiments,
-        height,
-        nbItemsPerRow,
-        completeColorScale,
-        new Set(colorScale?.domain)
-      )
+      getCustomPlotData(value, experiments, height, nbItemsPerRow, colorScale)
     )
   }