Add Custom Plots Section

extension/src/experiments/index.ts

@@ -438,6 +438,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.experiments.getFinishedExperiments()
   }
 
+  public getExperiments(): Experiment[] {
+    return this.experiments.getExperiments()
+  }
+
   public getExperimentDisplayName(experimentId: string) {
     const experiment = this.experiments
       .getCombinedList()
@@ -500,6 +504,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/experiments/model/index.ts

@@ -351,7 +351,7 @@ export class ExperimentsModel extends ModelWithPersistence {
     return collectFlatExperimentParams(params)
   }
 
-  public getExperiments() {
+  public getExperiments(): Experiment[] {
     return this.getExperimentsAndQueued().filter(({ status }) => {
       return !isQueued(status)
     })

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,4 +1,5 @@
 import omit from 'lodash.omit'
+import get from 'lodash.get'
 import { TopLevelSpec } from 'vega-lite'
 import { VisualizationSpec } from 'react-vega'
 import { getRevisionFirstThreeColumns } from './util'
@@ -13,7 +14,8 @@ import {
   TemplatePlotEntry,
   TemplatePlotSection,
   PlotsType,
-  Revision
+  Revision,
+  CustomPlotData
 } from '../webview/contract'
 import {
   EXPERIMENT_WORKSPACE_ID,
@@ -28,9 +30,11 @@ import {
 import { extractColumns } from '../../experiments/columns/extract'
 import {
   decodeColumn,
-  appendColumnToPath
+  appendColumnToPath,
+  splitColumnPath
 } from '../../experiments/columns/paths'
 import {
+  ColumnType,
   Experiment,
   isRunning,
   MetricOrParamColumns
@@ -243,6 +247,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: { metric: string; param: string }[],
+  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,
@@ -45,11 +48,13 @@ import {
   MultiSourceVariations
 } from '../multiSource/collect'
 import { isDvcError } from '../../cli/dvc/reader'
+import { Toast } from '../../vscode/toast'
 
 export class PlotsModel extends ModelWithPersistence {
   private readonly experiments: Experiments
 
   private plotSizes: Record<Section, number>
+  private customPlotsOrder: { metric: string; param: string }[]
   private sectionCollapsed: SectionCollapsed
   private commitRevisions: Record<string, string> = {}
 
@@ -64,6 +69,7 @@ export class PlotsModel extends ModelWithPersistence {
   private multiSourceEncoding: MultiSourceEncoding = {}
 
   private checkpointPlots?: CheckpointPlot[]
+  private customPlots?: CustomPlotData[]
   private selectedMetrics?: string[]
   private metricOrder: string[]
 
@@ -89,6 +95,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) {
@@ -119,6 +127,8 @@ export class PlotsModel extends ModelWithPersistence {
         collectMultiSourceVariations(data, this.multiSourceVariations)
       ])
 
+    this.recreateCustomPlots()
+
     this.comparisonData = {
       ...this.comparisonData,
       ...comparisonData
@@ -127,7 +137,6 @@ export class PlotsModel extends ModelWithPersistence {
       ...this.revisionData,
       ...revisionData
     }
-
     this.templates = { ...this.templates, ...templates }
     this.multiSourceVariations = multiSourceVariations
     this.multiSourceEncoding = collectMultiSourceEncoding(
@@ -171,6 +180,58 @@ 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: { metric: string; param: string }[]) {
+    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: { metric: string; param: string }) {
+    const plotAlreadyExists = this.getCustomPlotsOrder().some(
+      ({ param, metric }) =>
+        param === metricAndParam.param && metric === metricAndParam.metric
+    )
+
+    if (plotAlreadyExists) {
+      return Toast.showError('Custom plot already exists.')
+    }
+
+    const newCustomPlotsOrder = [...this.getCustomPlotsOrder(), metricAndParam]
+    this.setCustomPlotsOrder(newCustomPlotsOrder)
+  }
+
   public setupManualRefresh(id: string) {
     this.deleteRevisionData(id)
   }

extension/src/plots/model/quickPick.ts

@@ -0,0 +1,65 @@
+import { getCustomPlotId } from './collect'
+import { splitColumnPath } from '../../experiments/columns/paths'
+import { pickFromColumnLikes } from '../../experiments/columns/quickPick'
+import { Column, ColumnType } from '../../experiments/webview/contract'
+import { definedAndNonEmpty } from '../../util/array'
+import {
+  quickPickManyValues,
+  QuickPickOptionsWithTitle
+} from '../../vscode/quickPick'
+import { Title } from '../../vscode/title'
+import { Toast } from '../../vscode/toast'
+
+export const pickCustomPlots = (
+  plots: { metric: string; param: string }[],
+  quickPickOptions: QuickPickOptionsWithTitle
+): Thenable<string[] | undefined> => {
+  if (!definedAndNonEmpty(plots)) {
+    return Toast.showError('There are no plots to remove.')
+  }
+
+  const plotsItems = plots.map(({ metric, param }) => {
+    const splitMetric = splitColumnPath(metric)
+    const splitParam = splitColumnPath(param)
+    return {
+      description: `${metric} vs ${param}`,
+      label: `${splitMetric[splitMetric.length - 1]} vs ${
+        splitParam[splitParam.length - 1]
+      }`,
+      value: getCustomPlotId(metric, param)
+    }
+  })
+
+  return quickPickManyValues(plotsItems, quickPickOptions)
+}
+
+export const pickMetricAndParam = async (columns: Column[]) => {
+  const metricColumnLikes = columns
+    .filter(({ type }) => type === ColumnType.METRICS)
+    .map(({ label, path }) => ({ label, path }))
+  const paramColumnLikes = columns
+    .filter(({ type }) => type === ColumnType.PARAMS)
+    .map(({ label, path }) => ({ label, path }))
+  if (
+    !definedAndNonEmpty(metricColumnLikes) ||
+    !definedAndNonEmpty(paramColumnLikes)
+  ) {
+    return Toast.showError('There are no metrics or params to select from.')
+  }
+  const metric = await pickFromColumnLikes(metricColumnLikes, {
+    title: Title.SELECT_METRIC_CUSTOM_PLOT
+  })
+
+  if (!metric) {
+    return
+  }
+
+  const param = await pickFromColumnLikes(paramColumnLikes, {
+    title: Title.SELECT_PARAM_CUSTOM_PLOT
+  })
+
+  if (!param) {
+    return
+  }
+  return { metric: metric.path, param: param.path }
+}