Apply selected experiment colors to custom plots (#4647)

Author: mattseddon

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

@@ -41,13 +41,31 @@ describe('collectCustomPlots', () => {
   it('should return the expected data from the test fixture', () => {
     const expectedOutput: CustomPlotData[] = customPlotsFixture.plots
     const data = collectCustomPlots({
+      colorScale: {
+        domain: ['main', 'exp-e7a67', 'test-branch', 'exp-83425'],
+        range: ['#13adc7', '#f46837', '#48bb78', '#4299e1']
+      },
       experiments: experimentsWithCommits,
       height: DEFAULT_PLOT_HEIGHT,
       nbItemsPerRow: DEFAULT_NB_ITEMS_PER_ROW,
       plotsOrderValues: customPlotsOrderFixture
     })
     expect(data).toStrictEqual(expectedOutput)
   })
+
+  it('should return selected revisions last', () => {
+    const data = collectCustomPlots({
+      colorScale: {
+        domain: ['main'],
+        range: ['#13adc7']
+      },
+      experiments: experimentsWithCommits,
+      height: DEFAULT_PLOT_HEIGHT,
+      nbItemsPerRow: DEFAULT_NB_ITEMS_PER_ROW,
+      plotsOrderValues: customPlotsOrderFixture
+    })
+    expect(data[0].values.slice(-1)[0].id).toStrictEqual('main')
+  })
 })
 
 describe('collectData', () => {

extension/src/plots/model/collect.ts

@@ -42,6 +42,7 @@ import {
   getParent,
   getPathArray
 } from '../../fileSystem/util'
+import { Color } from '../../experiments/model/status/colors'
 
 export const getCustomPlotId = (metric: string, param: string) =>
   `custom-${metric}-${param}`
@@ -52,21 +53,28 @@ const getValueFromColumn = (path: string, experiment: Experiment) =>
 const getValues = (
   experiments: Experiment[],
   metricPath: string,
-  paramPath: string
+  paramPath: string,
+  renderLastIds: Set<string> = new Set<string>()
 ): CustomPlotValues => {
   const values: CustomPlotValues = []
 
   for (const experiment of experiments) {
     const metricValue = getValueFromColumn(metricPath, experiment)
     const paramValue = getValueFromColumn(paramPath, experiment)
 
-    if (metricValue !== undefined && paramValue !== undefined) {
-      values.push({
-        id: experiment.id,
-        metric: metricValue,
-        param: paramValue
-      })
+    if (metricValue === undefined || paramValue === undefined) {
+      continue
     }
+
+    const value = {
+      id: experiment.id,
+      metric: metricValue,
+      param: paramValue
+    }
+
+    renderLastIds.has(experiment.id)
+      ? values.push(value)
+      : values.unshift(value)
   }
 
   return values
@@ -76,13 +84,15 @@ const getCustomPlotData = (
   orderValue: CustomPlotsOrderValue,
   experiments: Experiment[],
   height: number,
-  nbItemsPerRow: number
+  nbItemsPerRow: number,
+  completeColorScale: ColorScale,
+  renderLastIds: Set<string>
 ): CustomPlotData => {
   const { metric, param } = orderValue
   const metricPath = getFullValuePath(ColumnType.METRICS, metric)
   const paramPath = getFullValuePath(ColumnType.PARAMS, param)
 
-  const values = getValues(experiments, metricPath, paramPath)
+  const values = getValues(experiments, metricPath, paramPath, renderLastIds)
 
   const [{ param: paramVal, metric: metricVal }] = values
   const yTitle = truncateVerticalTitle(metric, nbItemsPerRow, height) as string
@@ -92,7 +102,8 @@ const getCustomPlotData = (
     metric,
     param,
     typeof metricVal,
-    typeof paramVal
+    typeof paramVal,
+    completeColorScale
   )
 
   return {
@@ -104,21 +115,54 @@ 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,
   experiments,
   height,
   nbItemsPerRow
 }: {
+  colorScale: ColorScale | undefined
   plotsOrderValues: CustomPlotsOrderValue[]
   experiments: Experiment[]
   height: number
   nbItemsPerRow: number
 }): CustomPlotData[] => {
   const plots = []
 
+  const completeColorScale = fillColorScale(colorScale, experiments)
+
   for (const value of plotsOrderValues) {
-    plots.push(getCustomPlotData(value, experiments, height, nbItemsPerRow))
+    plots.push(
+      getCustomPlotData(
+        value,
+        experiments,
+        height,
+        nbItemsPerRow,
+        completeColorScale,
+        new Set(colorScale?.domain)
+      )
+    )
   }
 
   return plots

extension/src/plots/model/custom.ts

@@ -2,6 +2,7 @@ import { VisualizationSpec } from 'react-vega'
 import { getCustomPlotId } from './collect'
 import { Column, ColumnType } from '../../experiments/webview/contract'
 import { FILE_SEPARATOR } from '../../experiments/columns/paths'
+import { ColorScale } from '../webview/contract'
 
 export type CustomPlotsOrderValue = {
   metric: string
@@ -77,12 +78,14 @@ export const createSpec = (
   metric: string,
   param: string,
   metricType: string,
-  paramType: string
+  paramType: string,
+  colorScale: ColorScale
 ) =>
   ({
     $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
     data: { name: 'values' },
     encoding: {
+      color: { field: 'id', legend: null, scale: colorScale },
       x: {
         axis: {
           labelLimit: 75,

extension/src/plots/model/index.ts

@@ -166,8 +166,14 @@ export class PlotsModel extends ModelWithPersistence {
     const nbItemsPerRow = this.getNbItemsPerRowOrWidth(
       PlotsSection.CUSTOM_PLOTS
     )
+    const colorScale = getColorScale(
+      this.getSelectedRevisionDetails().filter(
+        ({ id }) => id !== EXPERIMENT_WORKSPACE_ID
+      )
+    )
     const plotsOrderValues = this.getCustomPlotsOrder()
     const plots: CustomPlotData[] = collectCustomPlots({
+      colorScale,
       experiments,
       height,
       nbItemsPerRow,

extension/src/test/fixtures/expShow/base/customPlots.ts

@@ -103,6 +103,28 @@ const data: CustomPlotsData = {
         $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
         data: { name: 'values' },
         encoding: {
+          color: {
+            field: 'id',
+            legend: null,
+            scale: {
+              domain: [
+                'main',
+                'exp-e7a67',
+                'test-branch',
+                'exp-83425',
+                'fe2919b',
+                '7df876c'
+              ],
+              range: [
+                '#13adc7',
+                '#f46837',
+                '#48bb78',
+                '#4299e1',
+                '#4c78a8',
+                '#4c78a8'
+              ]
+            }
+          },
           x: {
             axis: {
               labelLimit: 75,
@@ -157,13 +179,13 @@ const data: CustomPlotsData = {
         width: 'container'
       },
       values: [
-        { id: 'main', metric: 2.048856019973755, param: 'logs.csv' },
+        { id: '7df876c', metric: 2.048856019973755, param: 'logs.csv' },
         {
           id: 'fe2919b',
           metric: 2.048856019973755,
           param: 'logs.csv'
         },
-        { id: '7df876c', metric: 2.048856019973755, param: 'logs.csv' },
+        { id: 'main', metric: 2.048856019973755, param: 'logs.csv' },
         { id: 'exp-e7a67', metric: 2.0205044746398926, param: 'logs.csv' },
         { id: 'test-branch', metric: 1.9293040037155151, param: 'logs.csv' },
         {
@@ -179,7 +201,7 @@ const data: CustomPlotsData = {
       param: 'params.yaml:epochs',
       values: [
         {
-          id: 'main',
+          id: '7df876c',
           metric: 0.3484833240509033,
           param: 5
         },
@@ -189,7 +211,7 @@ const data: CustomPlotsData = {
           param: 5
         },
         {
-          id: '7df876c',
+          id: 'main',
           metric: 0.3484833240509033,
           param: 5
         },
@@ -213,6 +235,28 @@ const data: CustomPlotsData = {
         $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
         data: { name: 'values' },
         encoding: {
+          color: {
+            field: 'id',
+            legend: null,
+            scale: {
+              domain: [
+                'main',
+                'exp-e7a67',
+                'test-branch',
+                'exp-83425',
+                'fe2919b',
+                '7df876c'
+              ],
+              range: [
+                '#13adc7',
+                '#f46837',
+                '#48bb78',
+                '#4299e1',
+                '#4c78a8',
+                '#4c78a8'
+              ]
+            }
+          },
           x: {
             axis: {
               labelLimit: 75,

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

@@ -490,12 +490,23 @@ suite('Plots Test Suite', () => {
         type: MessageFromWebviewType.EXPORT_PLOT_DATA_AS_JSON
       })
 
+      const expectedOrder = [
+        'exp-83425',
+        'test-branch',
+        'exp-e7a67',
+        '7df876c',
+        'fe2919b',
+        'main'
+      ]
+
       await openFileEvent
 
       expect(mockWriteJson).to.be.calledOnce
       expect(mockWriteJson).to.be.calledWithExactly(
         exportFile.path,
-        customPlot.values,
+        [...customPlot.values].sort(
+          (a, b) => expectedOrder.indexOf(a.id) - expectedOrder.indexOf(b.id)
+        ),
         true
       )
       expect(mockOpenFile).to.calledWithExactly(exportFile.path)
@@ -579,11 +590,23 @@ suite('Plots Test Suite', () => {
 
       await openFileEvent
 
+      const expectedOrder = [
+        'exp-83425',
+        'test-branch',
+        'exp-e7a67',
+        '7df876c',
+        'fe2919b',
+        'main'
+      ]
+
       expect(mockWriteTsv).to.be.calledOnce
       expect(mockWriteTsv).to.be.calledWithExactly(
         exportFile.path,
-        customPlot.values
+        [...customPlot.values].sort(
+          (a, b) => expectedOrder.indexOf(a.id) - expectedOrder.indexOf(b.id)
+        )
       )
+
       expect(mockOpenFile).to.calledWithExactly(exportFile.path)
       expect(mockSendTelemetryEvent).to.be.calledOnce
       expect(mockSendTelemetryEvent).to.be.calledWithExactly(