Add Custom Plots Section (#3342)

Author: julieg18

extension/src/experiments/index.ts

@@ -439,6 +439,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.experiments.getFinishedExperiments()
   }
 
+  public getExperiments() {
+    return this.experiments.getExperiments()
+  }
+
   public getExperimentDisplayName(experimentId: string) {
     const experiment = this.experiments
       .getCombinedList()
@@ -501,6 +505,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.columns.getFirstThreeColumnOrder()
   }
 
+  public getColumnTerminalNodes() {
+    return this.columns.getTerminalNodes()
+  }
+
   public getHasData() {
     if (this.deferred.state === 'none') {
       return

extension/src/persistence/constants.ts

@@ -10,6 +10,7 @@ export enum PersistenceKey {
   PLOT_COMPARISON_ORDER = 'plotComparisonOrder:',
   PLOT_COMPARISON_PATHS_ORDER = 'plotComparisonPathsOrder',
   PLOT_METRIC_ORDER = 'plotMetricOrder:',
+  PLOTS_CUSTOM_ORDER = 'plotCustomOrder:',
   PLOT_SECTION_COLLAPSED = 'plotSectionCollapsed:',
   PLOT_SELECTED_METRICS = 'plotSelectedMetrics:',
   PLOT_SIZES = 'plotSizes:',

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

@@ -6,12 +6,14 @@ import {
   collectCheckpointPlotsData,
   collectTemplates,
   collectMetricOrder,
-  collectOverrideRevisionDetails
+  collectOverrideRevisionDetails,
+  collectCustomPlotsData
 } from './collect'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
 import expShowFixture from '../../test/fixtures/expShow/base/output'
 import modifiedFixture from '../../test/fixtures/expShow/modified/output'
 import checkpointPlotsFixture from '../../test/fixtures/expShow/base/checkpointPlots'
+import customPlotsFixture from '../../test/fixtures/expShow/base/customPlots'
 import {
   ExperimentsOutput,
   ExperimentStatus,
@@ -27,6 +29,62 @@ const logsLossPath = join('logs', 'loss.tsv')
 
 const logsLossPlot = (plotsDiffFixture[logsLossPath][0] || {}) as TemplatePlot
 
+describe('collectCustomPlotsData', () => {
+  it('should return the expected data from the text fixture', () => {
+    const data = collectCustomPlotsData(
+      [
+        {
+          metric: 'metrics:summary.json:loss',
+          param: 'params:params.yaml:dropout'
+        },
+        {
+          metric: 'metrics:summary.json:accuracy',
+          param: 'params:params.yaml:epochs'
+        }
+      ],
+      [
+        {
+          id: '12345',
+          label: '123',
+          metrics: {
+            'summary.json': {
+              accuracy: 0.4668000042438507,
+              loss: 2.0205044746398926
+            }
+          },
+          name: 'exp-e7a67',
+          params: { 'params.yaml': { dropout: 0.15, epochs: 16 } }
+        },
+        {
+          id: '12345',
+          label: '123',
+          metrics: {
+            'summary.json': {
+              accuracy: 0.3484833240509033,
+              loss: 1.9293040037155151
+            }
+          },
+          name: 'exp-83425',
+          params: { 'params.yaml': { dropout: 0.25, epochs: 10 } }
+        },
+        {
+          id: '12345',
+          label: '123',
+          metrics: {
+            'summary.json': {
+              accuracy: 0.6768440509033,
+              loss: 2.298503875732422
+            }
+          },
+          name: 'exp-f13bca',
+          params: { 'params.yaml': { dropout: 0.32, epochs: 20 } }
+        }
+      ]
+    )
+    expect(data).toStrictEqual(customPlotsFixture.plots)
+  })
+})
+
 describe('collectCheckpointPlotsData', () => {
   it('should return the expected data from the test fixture', () => {
     const data = collectCheckpointPlotsData(expShowFixture)

extension/src/plots/model/collect.ts

@@ -1,6 +1,8 @@
 import omit from 'lodash.omit'
+import get from 'lodash.get'
 import { TopLevelSpec } from 'vega-lite'
 import { VisualizationSpec } from 'react-vega'
+import { CustomPlotsOrderValue } from '.'
 import { getRevisionFirstThreeColumns } from './util'
 import {
   ColorScale,
@@ -13,7 +15,8 @@ import {
   TemplatePlotEntry,
   TemplatePlotSection,
   PlotsType,
-  Revision
+  Revision,
+  CustomPlotData
 } from '../webview/contract'
 import {
   EXPERIMENT_WORKSPACE_ID,
@@ -28,9 +31,11 @@ import {
 import { extractColumns } from '../../experiments/columns/extract'
 import {
   decodeColumn,
-  appendColumnToPath
+  appendColumnToPath,
+  splitColumnPath
 } from '../../experiments/columns/paths'
 import {
+  ColumnType,
   Experiment,
   isRunning,
   MetricOrParamColumns
@@ -243,6 +248,48 @@ export const collectCheckpointPlotsData = (
   return plotsData
 }
 
+export const getCustomPlotId = (metric: string, param: string) =>
+  `custom-${metric}-${param}`
+
+const collectCustomPlotData = (
+  metric: string,
+  param: string,
+  experiments: Experiment[]
+): CustomPlotData => {
+  const splitUpMetricPath = splitColumnPath(metric)
+  const splitUpParamPath = splitColumnPath(param)
+  const plotData: CustomPlotData = {
+    id: getCustomPlotId(metric, param),
+    metric: metric.slice(ColumnType.METRICS.length + 1),
+    param: param.slice(ColumnType.PARAMS.length + 1),
+    values: []
+  }
+
+  for (const experiment of experiments) {
+    const metricValue = get(experiment, splitUpMetricPath) as number | undefined
+    const paramValue = get(experiment, splitUpParamPath) as number | undefined
+
+    if (metricValue !== undefined && paramValue !== undefined) {
+      plotData.values.push({
+        expName: experiment.name || experiment.label,
+        metric: metricValue,
+        param: paramValue
+      })
+    }
+  }
+
+  return plotData
+}
+
+export const collectCustomPlotsData = (
+  metricsAndParams: CustomPlotsOrderValue[],
+  experiments: Experiment[]
+): CustomPlotData[] => {
+  return metricsAndParams.map(({ metric, param }) =>
+    collectCustomPlotData(metric, param, experiments)
+  )
+}
+
 type MetricOrderAccumulator = {
   newOrder: string[]
   uncollectedMetrics: string[]

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

@@ -103,6 +103,7 @@ describe('plotsModel', () => {
     const expectedSectionCollapsed = {
       [Section.CHECKPOINT_PLOTS]: true,
       [Section.TEMPLATE_PLOTS]: false,
+      [Section.CUSTOM_PLOTS]: false,
       [Section.COMPARISON_TABLE]: false
     }
 

extension/src/plots/model/index.ts

@@ -11,7 +11,9 @@ import {
   RevisionData,
   TemplateAccumulator,
   collectCommitRevisionDetails,
-  collectOverrideRevisionDetails
+  collectOverrideRevisionDetails,
+  collectCustomPlotsData,
+  getCustomPlotId
 } from './collect'
 import { getRevisionFirstThreeColumns } from './util'
 import {
@@ -24,7 +26,8 @@ import {
   DEFAULT_SECTION_SIZES,
   Section,
   SectionCollapsed,
-  PlotSizeNumber
+  PlotSizeNumber,
+  CustomPlotData
 } from '../webview/contract'
 import {
   ExperimentsOutput,
@@ -46,10 +49,13 @@ import {
 } from '../multiSource/collect'
 import { isDvcError } from '../../cli/dvc/reader'
 
+export type CustomPlotsOrderValue = { metric: string; param: string }
+
 export class PlotsModel extends ModelWithPersistence {
   private readonly experiments: Experiments
 
   private plotSizes: Record<Section, number>
+  private customPlotsOrder: CustomPlotsOrderValue[]
   private sectionCollapsed: SectionCollapsed
   private commitRevisions: Record<string, string> = {}
 
@@ -64,6 +70,7 @@ export class PlotsModel extends ModelWithPersistence {
   private multiSourceEncoding: MultiSourceEncoding = {}
 
   private checkpointPlots?: CheckpointPlot[]
+  private customPlots?: CustomPlotData[]
   private selectedMetrics?: string[]
   private metricOrder: string[]
 
@@ -89,6 +96,8 @@ export class PlotsModel extends ModelWithPersistence {
       undefined
     )
     this.metricOrder = this.revive(PersistenceKey.PLOT_METRIC_ORDER, [])
+
+    this.customPlotsOrder = this.revive(PersistenceKey.PLOTS_CUSTOM_ORDER, [])
   }
 
   public transformAndSetExperiments(data: ExperimentsOutput) {
@@ -102,6 +111,8 @@ export class PlotsModel extends ModelWithPersistence {
 
     this.setMetricOrder()
 
+    this.recreateCustomPlots()
+
     return this.removeStaleData()
   }
 
@@ -119,6 +130,8 @@ export class PlotsModel extends ModelWithPersistence {
         collectMultiSourceVariations(data, this.multiSourceVariations)
       ])
 
+    this.recreateCustomPlots()
+
     this.comparisonData = {
       ...this.comparisonData,
       ...comparisonData
@@ -127,7 +140,6 @@ export class PlotsModel extends ModelWithPersistence {
       ...this.revisionData,
       ...revisionData
     }
-
     this.templates = { ...this.templates, ...templates }
     this.multiSourceVariations = multiSourceVariations
     this.multiSourceEncoding = collectMultiSourceEncoding(
@@ -171,6 +183,49 @@ export class PlotsModel extends ModelWithPersistence {
     }
   }
 
+  public getCustomPlots() {
+    if (!this.customPlots) {
+      return
+    }
+    return {
+      plots: this.customPlots,
+      size: this.getPlotSize(Section.CUSTOM_PLOTS)
+    }
+  }
+
+  public recreateCustomPlots() {
+    const customPlots: CustomPlotData[] = collectCustomPlotsData(
+      this.getCustomPlotsOrder(),
+      this.experiments.getExperiments()
+    )
+    this.customPlots = customPlots
+  }
+
+  public getCustomPlotsOrder() {
+    return this.customPlotsOrder
+  }
+
+  public setCustomPlotsOrder(plotsOrder: CustomPlotsOrderValue[]) {
+    this.customPlotsOrder = plotsOrder
+    this.persist(PersistenceKey.PLOTS_CUSTOM_ORDER, this.customPlotsOrder)
+    this.recreateCustomPlots()
+  }
+
+  public removeCustomPlots(plotIds: string[]) {
+    const newCustomPlotsOrder = this.getCustomPlotsOrder().filter(
+      ({ metric, param }) => {
+        return !plotIds.includes(getCustomPlotId(metric, param))
+      }
+    )
+
+    this.setCustomPlotsOrder(newCustomPlotsOrder)
+  }
+
+  public addCustomPlot(metricAndParam: CustomPlotsOrderValue) {
+    const newCustomPlotsOrder = [...this.getCustomPlotsOrder(), metricAndParam]
+    this.setCustomPlotsOrder(newCustomPlotsOrder)
+  }
+
   public setupManualRefresh(id: string) {
     this.deleteRevisionData(id)
   }