Support user_attrs on GraphHistory

optuna_dashboard/ts/components/GraphHistory.tsx

@@ -1,5 +1,5 @@
 import * as plotly from "plotly.js-dist-min"
-import React, { ChangeEvent, FC, useEffect, useState } from "react"
+import React, { ChangeEvent, FC, useEffect, useMemo, useState } from "react"
 import {
   Grid,
   FormControl,
@@ -19,19 +19,128 @@ import { plotlyDarkTemplate } from "./PlotlyDarkMode"
 
 const plotDomId = "graph-history"
 
+class Target {
+  kind: "objective" | "user_attr"
+  key: number | string
+
+  constructor(kind: "objective" | "user_attr", key: number | string) {
+    this.kind = kind
+    this.key = key
+  }
+
+  validate(): boolean {
+    if (this.kind === "objective") {
+      if (typeof this.key !== "number") {
+        return false
+      }
+    } else if (this.kind === "user_attr") {
+      if (typeof this.key !== "string") {
+        return false
+      }
+    } else {
+      return false
+    }
+    return true
+  }
+
+  toLabel(objectiveNames: string[]): string {
+    if (this.kind === "objective") {
+      const objectiveId: number = this.key as number
+      if (objectiveNames.length > objectiveId) {
+        return objectiveNames[objectiveId]
+      }
+      return `Objective ${objectiveId}`
+    } else {
+      return `User Attribute ${this.key}`
+    }
+  }
+
+  getObjectiveId(): number | null {
+    return this.key as number
+  }
+
+  getTargetValue(trial: Trial): number | null {
+    if (!this.validate()) {
+      return null
+    }
+    if (this.kind === "objective") {
+      const objectiveId = this.getObjectiveId()
+      if (
+        objectiveId === null ||
+        trial.values === undefined ||
+        trial.values.length <= objectiveId
+      ) {
+        return null
+      }
+      const value = trial.values[objectiveId]
+      if (value === "inf" || value === "-inf") {
+        return null
+      }
+      return value
+    } else if (this.kind === "user_attr") {
+      const attr = trial.user_attrs.find((attr) => attr.key === this.key)
+      if (attr === undefined) {
+        return null
+      }
+      const value = Number(attr.value)
+      if (value === undefined) {
+        return null
+      }
+      return value
+    }
+    return null
+  }
+}
+
 export const GraphHistory: FC<{
   study: StudyDetail | null
 }> = ({ study = null }) => {
   const theme = useTheme()
   const [xAxis, setXAxis] = useState<string>("number")
-  const [objectiveId, setObjectiveId] = useState<number>(0)
+  const [targetIndex, setTargetIndex] = useState<number>(0)
   const [logScale, setLogScale] = useState<boolean>(false)
   const [filterCompleteTrial, setFilterCompleteTrial] = useState<boolean>(false)
   const [filterPrunedTrial, setFilterPrunedTrial] = useState<boolean>(false)
+  const [targetList, setTargetList] = useState<Target[]>([])
   const objectiveNames: string[] = study?.objective_names || []
 
+  useMemo(() => {
+    if (study !== null) {
+      const targets: Target[] = [
+        ...study.directions.map((v, i) => new Target("objective", i)),
+        ...study.union_user_attrs
+          .filter((attr) => attr.sortable)
+          .map((attr) => new Target("user_attr", attr.key)),
+      ]
+      setTargetList(targets)
+    }
+  }, [study?.directions, study?.union_user_attrs])
+
+  useEffect(() => {
+    if (study !== null) {
+      plotHistory(
+        study,
+        targetList[targetIndex],
+        xAxis,
+        logScale,
+        filterCompleteTrial,
+        filterPrunedTrial,
+        theme.palette.mode
+      )
+    }
+  }, [
+    study,
+    targetIndex,
+    targetList,
+    logScale,
+    xAxis,
+    filterPrunedTrial,
+    filterCompleteTrial,
+    theme.palette.mode,
+  ])
+
   const handleObjectiveChange = (event: SelectChangeEvent<number>) => {
-    setObjectiveId(event.target.value as number)
+    setTargetIndex(event.target.value as number)
   }
 
   const handleXAxisChange = (e: ChangeEvent<HTMLInputElement>) => {
@@ -50,28 +159,6 @@ export const GraphHistory: FC<{
     setFilterPrunedTrial(!filterPrunedTrial)
   }
 
-  useEffect(() => {
-    if (study !== null) {
-      plotHistory(
-        study,
-        objectiveId,
-        xAxis,
-        logScale,
-        filterCompleteTrial,
-        filterPrunedTrial,
-        theme.palette.mode
-      )
-    }
-  }, [
-    study,
-    objectiveId,
-    logScale,
-    xAxis,
-    filterPrunedTrial,
-    filterCompleteTrial,
-    theme.palette.mode,
-  ])
-
   return (
     <Grid container direction="row">
       <Grid
@@ -84,18 +171,16 @@ export const GraphHistory: FC<{
         <Typography variant="h6" sx={{ margin: "1em 0", fontWeight: 600 }}>
           History
         </Typography>
-        {study !== null && study.directions.length !== 1 ? (
+        {study !== null && targetList.length >= 2 ? (
           <FormControl
             component="fieldset"
             sx={{ marginBottom: theme.spacing(2) }}
           >
-            <FormLabel component="legend">Objective ID:</FormLabel>
-            <Select value={objectiveId} onChange={handleObjectiveChange}>
-              {study.directions.map((d, i) => (
+            <FormLabel component="legend">y Axis</FormLabel>
+            <Select value={targetIndex} onChange={handleObjectiveChange}>
+              {targetList.map((t, i) => (
                 <MenuItem value={i} key={i}>
-                  {objectiveNames.length === study?.directions.length
-                    ? objectiveNames[i]
-                    : `${i}`}
+                  {t.toLabel(objectiveNames)}
                 </MenuItem>
               ))}
             </Select>
@@ -173,23 +258,17 @@ export const GraphHistory: FC<{
   )
 }
 
-const filterFunc = (trial: Trial, objectiveId: number): boolean => {
+const filterFunc = (trial: Trial, target: Target): boolean => {
   if (trial.state !== "Complete" && trial.state !== "Pruned") {
     return false
   }
-  if (trial.values === undefined) {
-    return false
-  }
-  return (
-    trial.values.length > objectiveId &&
-    trial.values[objectiveId] !== "inf" &&
-    trial.values[objectiveId] !== "-inf"
-  )
+  const value = target.getTargetValue(trial)
+  return value !== null
 }
 
 const plotHistory = (
   study: StudyDetail,
-  objectiveId: number,
+  target: Target,
   xAxis: string,
   logScale: boolean,
   filterCompleteTrial: boolean,
@@ -219,7 +298,7 @@ const plotHistory = (
     template: mode === "dark" ? plotlyDarkTemplate : {},
   }
 
-  let filteredTrials = study.trials.filter((t) => filterFunc(t, objectiveId))
+  let filteredTrials = study.trials.filter((t) => filterFunc(t, target))
   if (filterCompleteTrial) {
     filteredTrials = filteredTrials.filter((t) => t.state !== "Complete")
   }
@@ -239,61 +318,64 @@ const plotHistory = (
       : trial.datetime_complete!
   }
 
-  const xForLinePlot: (number | Date)[] = []
-  const yForLinePlot: number[] = []
-  let currentBest: number | null = null
-  for (let i = 0; i < filteredTrials.length; i++) {
-    const t = filteredTrials[i]
-    if (currentBest === null) {
-      currentBest = t.values![objectiveId] as number
-      xForLinePlot.push(getAxisX(t))
-      yForLinePlot.push(t.values![objectiveId] as number)
-    } else if (
-      study.directions[objectiveId] === "maximize" &&
-      t.values![objectiveId] > currentBest
-    ) {
-      const p = filteredTrials[i - 1]
-      if (!xForLinePlot.includes(getAxisX(p))) {
-        xForLinePlot.push(getAxisX(p))
-        yForLinePlot.push(currentBest)
-      }
-      currentBest = t.values![objectiveId] as number
-      xForLinePlot.push(getAxisX(t))
-      yForLinePlot.push(t.values![objectiveId] as number)
-    } else if (
-      study.directions[objectiveId] === "minimize" &&
-      t.values![objectiveId] < currentBest
-    ) {
-      const p = filteredTrials[i - 1]
-      if (!xForLinePlot.includes(getAxisX(p))) {
-        xForLinePlot.push(getAxisX(p))
-        yForLinePlot.push(currentBest)
-      }
-      currentBest = t.values![objectiveId] as number
-      xForLinePlot.push(getAxisX(t))
-      yForLinePlot.push(t.values![objectiveId] as number)
-    }
-  }
-  xForLinePlot.push(getAxisX(filteredTrials[filteredTrials.length - 1]))
-  yForLinePlot.push(yForLinePlot[yForLinePlot.length - 1])
-
   const plotData: Partial<plotly.PlotData>[] = [
     {
       x: filteredTrials.map(getAxisX),
       y: filteredTrials.map(
-        (t: Trial): number => t.values![objectiveId] as number
+        (t: Trial): number => target.getTargetValue(t) as number
       ),
       name: "Objective Value",
       mode: "markers",
       type: "scatter",
     },
-    {
+  ]
+
+  const objectiveId = target.getObjectiveId()
+  if (objectiveId !== null) {
+    const xForLinePlot: (number | Date)[] = []
+    const yForLinePlot: number[] = []
+    let currentBest: number | null = null
+    for (let i = 0; i < filteredTrials.length; i++) {
+      const t = filteredTrials[i]
+      if (currentBest === null) {
+        currentBest = t.values![objectiveId] as number
+        xForLinePlot.push(getAxisX(t))
+        yForLinePlot.push(t.values![objectiveId] as number)
+      } else if (
+        study.directions[objectiveId] === "maximize" &&
+        t.values![objectiveId] > currentBest
+      ) {
+        const p = filteredTrials[i - 1]
+        if (!xForLinePlot.includes(getAxisX(p))) {
+          xForLinePlot.push(getAxisX(p))
+          yForLinePlot.push(currentBest)
+        }
+        currentBest = t.values![objectiveId] as number
+        xForLinePlot.push(getAxisX(t))
+        yForLinePlot.push(t.values![objectiveId] as number)
+      } else if (
+        study.directions[objectiveId] === "minimize" &&
+        t.values![objectiveId] < currentBest
+      ) {
+        const p = filteredTrials[i - 1]
+        if (!xForLinePlot.includes(getAxisX(p))) {
+          xForLinePlot.push(getAxisX(p))
+          yForLinePlot.push(currentBest)
+        }
+        currentBest = t.values![objectiveId] as number
+        xForLinePlot.push(getAxisX(t))
+        yForLinePlot.push(t.values![objectiveId] as number)
+      }
+    }
+    xForLinePlot.push(getAxisX(filteredTrials[filteredTrials.length - 1]))
+    yForLinePlot.push(yForLinePlot[yForLinePlot.length - 1])
+    plotData.push({
       x: xForLinePlot,
       y: yForLinePlot,
       name: "Best Value",
       mode: "lines",
       type: "scatter",
-    },
-  ]
+    })
+  }
   plotly.react(plotDomId, plotData, layout)
 }