Refactor history component

optuna_dashboard/ts/components/GraphHistory.tsx

@@ -1,5 +1,5 @@
 import * as plotly from "plotly.js-dist-min"
-import React, { ChangeEvent, FC, useEffect, useMemo, useState } from "react"
+import React, { ChangeEvent, FC, useEffect, useState } from "react"
 import {
   Grid,
   FormControl,
@@ -16,120 +16,43 @@ import {
   useTheme,
 } from "@mui/material"
 import { plotlyDarkTemplate } from "./PlotlyDarkMode"
+import { useFilteredTrials, Target, useTargetList } from "../trialFilter"
 
 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 [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])
+  const objectiveNames: string[] = study?.objective_names || []
+  const targetList = useTargetList(study)
+  const [targetIndex, setTargetIndex] = useState<number>(0)
+  const trials = useFilteredTrials(
+    study,
+    targetList[targetIndex],
+    filterCompleteTrial,
+    filterPrunedTrial
+  )
 
   useEffect(() => {
     if (study !== null) {
       plotHistory(
-        study,
+        trials,
+        study.directions,
         targetList[targetIndex],
         xAxis,
         logScale,
-        filterCompleteTrial,
-        filterPrunedTrial,
         theme.palette.mode
       )
     }
   }, [
-    study,
+    trials,
+    study?.directions,
     targetIndex,
     targetList,
     logScale,
@@ -258,21 +181,12 @@ export const GraphHistory: FC<{
   )
 }
 
-const filterFunc = (trial: Trial, target: Target): boolean => {
-  if (trial.state !== "Complete" && trial.state !== "Pruned") {
-    return false
-  }
-  const value = target.getTargetValue(trial)
-  return value !== null
-}
-
 const plotHistory = (
-  study: StudyDetail,
+  trials: Trial[],
+  directions: StudyDirection[],
   target: Target,
   xAxis: string,
   logScale: boolean,
-  filterCompleteTrial: boolean,
-  filterPrunedTrial: boolean,
   mode: string
 ) => {
   if (document.getElementById(plotDomId) === null) {
@@ -297,15 +211,7 @@ const plotHistory = (
     showlegend: true,
     template: mode === "dark" ? plotlyDarkTemplate : {},
   }
-
-  let filteredTrials = study.trials.filter((t) => filterFunc(t, target))
-  if (filterCompleteTrial) {
-    filteredTrials = filteredTrials.filter((t) => t.state !== "Complete")
-  }
-  if (filterPrunedTrial) {
-    filteredTrials = filteredTrials.filter((t) => t.state !== "Pruned")
-  }
-  if (filteredTrials.length === 0) {
+  if (trials.length === 0) {
     plotly.react(plotDomId, [], layout)
     return
   }
@@ -320,10 +226,8 @@ const plotHistory = (
 
   const plotData: Partial<plotly.PlotData>[] = [
     {
-      x: filteredTrials.map(getAxisX),
-      y: filteredTrials.map(
-        (t: Trial): number => target.getTargetValue(t) as number
-      ),
+      x: trials.map(getAxisX),
+      y: trials.map((t: Trial): number => target.getTargetValue(t) as number),
       name: "Objective Value",
       mode: "markers",
       type: "scatter",
@@ -335,17 +239,17 @@ const plotHistory = (
     const xForLinePlot: (number | Date)[] = []
     const yForLinePlot: number[] = []
     let currentBest: number | null = null
-    for (let i = 0; i < filteredTrials.length; i++) {
-      const t = filteredTrials[i]
+    for (let i = 0; i < trials.length; i++) {
+      const t = trials[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" &&
+        directions[objectiveId] === "maximize" &&
         t.values![objectiveId] > currentBest
       ) {
-        const p = filteredTrials[i - 1]
+        const p = trials[i - 1]
         if (!xForLinePlot.includes(getAxisX(p))) {
           xForLinePlot.push(getAxisX(p))
           yForLinePlot.push(currentBest)
@@ -354,10 +258,10 @@ const plotHistory = (
         xForLinePlot.push(getAxisX(t))
         yForLinePlot.push(t.values![objectiveId] as number)
       } else if (
-        study.directions[objectiveId] === "minimize" &&
+        directions[objectiveId] === "minimize" &&
         t.values![objectiveId] < currentBest
       ) {
-        const p = filteredTrials[i - 1]
+        const p = trials[i - 1]
         if (!xForLinePlot.includes(getAxisX(p))) {
           xForLinePlot.push(getAxisX(p))
           yForLinePlot.push(currentBest)
@@ -367,7 +271,7 @@ const plotHistory = (
         yForLinePlot.push(t.values![objectiveId] as number)
       }
     }
-    xForLinePlot.push(getAxisX(filteredTrials[filteredTrials.length - 1]))
+    xForLinePlot.push(getAxisX(trials[trials.length - 1]))
     yForLinePlot.push(yForLinePlot[yForLinePlot.length - 1])
     plotData.push({
       x: xForLinePlot,

optuna_dashboard/ts/trialFilter.ts

@@ -0,0 +1,112 @@
+import { useMemo } from "react"
+
+export 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 useFilteredTrials = (
+  study: StudyDetail | null,
+  target: Target,
+  filterComplete: boolean,
+  filterPruned: boolean
+): Trial[] =>
+  useMemo<Trial[]>(() => {
+    if (study === null) {
+      return []
+    }
+    return study.trials.filter((t) => {
+      if (t.state !== "Complete" && t.state !== "Pruned") {
+        return false
+      }
+      if (t.state === "Complete" && filterComplete) {
+        return false
+      }
+      if (t.state === "Pruned" && filterPruned) {
+        return false
+      }
+      return target.getTargetValue(t) !== null
+    })
+  }, [study?.trials, target, filterComplete, filterPruned])
+
+export const useTargetList = (study: StudyDetail | null): Target[] =>
+  useMemo<Target[]>(() => {
+    if (study !== null) {
+      return [
+        ...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)),
+      ]
+    } else {
+      return [new Target("objective", 0)]
+    }
+  }, [study?.directions, study?.union_user_attrs])