From 270c406425c983468218f7fb790c599dc8b43a87 Mon Sep 17 00:00:00 2001
From: keisuke-umezawa <keisuke.umezawa@gmail.com>
Date: Thu, 15 Aug 2024 14:07:10 +0900
Subject: [PATCH] Add plotly color theme to PlotImportance

---
 tslib/react/src/components/PlotEdf.tsx     |   2 +-
 tslib/react/src/components/PlotHistory.tsx | 152 +++++++++------------
 2 files changed, 65 insertions(+), 89 deletions(-)

diff --git a/tslib/react/src/components/PlotEdf.tsx b/tslib/react/src/components/PlotEdf.tsx
index f1cc5aa7..65281dcf 100644
--- a/tslib/react/src/components/PlotEdf.tsx
+++ b/tslib/react/src/components/PlotEdf.tsx
@@ -4,7 +4,7 @@ import * as plotly from "plotly.js-dist-min"
 import { FC, useEffect, useMemo } from "react"
 import { useGraphComponentState } from "../hooks/useGraphComponentState"
 import { Target, useFilteredTrialsFromStudies } from "../utils/trialFilter"
-import { GraphContainer } from"./GraphContainer"
+import { GraphContainer } from "./GraphContainer"
 import { plotlyDarkTemplate } from "./PlotlyDarkMode"
 
 export type EdfPlotInfo = {
diff --git a/tslib/react/src/components/PlotHistory.tsx b/tslib/react/src/components/PlotHistory.tsx
index ed0728a8..6e235997 100644
--- a/tslib/react/src/components/PlotHistory.tsx
+++ b/tslib/react/src/components/PlotHistory.tsx
@@ -1,5 +1,4 @@
 import {
-  Checkbox,
   FormControl,
   FormControlLabel,
   FormLabel,
@@ -9,16 +8,21 @@ import {
   RadioGroup,
   Select,
   SelectChangeEvent,
-  Switch,
+  Slider,
   Typography,
   useTheme,
 } from "@mui/material"
 import * as Optuna from "@optuna/types"
 import * as plotly from "plotly.js-dist-min"
-import { ChangeEvent, FC, useEffect, useState, useMemo } from "react"
+import { ChangeEvent, FC, useEffect, useState } from "react"
 
 import { useGraphComponentState } from "../hooks/useGraphComponentState"
-import { Target, useFilteredTrialsFromStudies } from "../utils/trialFilter"
+import {
+  Target,
+  useFilteredTrialsFromStudies,
+  useObjectiveAndUserAttrTargetsFromStudies,
+} from "../utils/trialFilter"
+import { plotlyDarkTemplate } from "./PlotlyDarkMode"
 
 const plotDomId = "plot-history"
 
@@ -31,29 +35,31 @@ interface HistoryPlotInfo {
 
 export const PlotHistory: FC<{
   studies: Optuna.Study[]
-}> = ({ studies }) => {
+  colorTheme?: Partial<Plotly.Template>
+}> = ({ studies, colorTheme }) => {
+  const logScale = false
+  const filterPrunedTrial = false
+
   const { graphComponentState, notifyGraphDidRender } = useGraphComponentState()
 
   const theme = useTheme()
+  const colorThemeUsed =
+    colorTheme ?? (theme.palette.mode === "dark" ? plotlyDarkTemplate : {})
 
   const [xAxis, setXAxis] = useState<
     "number" | "datetime_start" | "datetime_complete"
   >("number")
-  
-  const [objectiveId, setObjectiveId] = useState<number>(0)
-  const [logScale, setLogScale] = useState<boolean>(false)
-  const [filterPrunedTrial, setFilterPrunedTrial] = useState<boolean>(false)
   const [markerSize, setMarkerSize] = useState<number>(5)
 
-  const target = useMemo<Target>(
-    () => new Target("objective", objectiveId),
-    [objectiveId]
-  )
+  const [targets, selected, setTarget] =
+    useObjectiveAndUserAttrTargetsFromStudies(studies)
+
   const trials = useFilteredTrialsFromStudies(
     studies,
-    [target],
-    filterPrunedTrial,
+    [selected],
+    filterPrunedTrial
   )
+
   const historyPlotInfos = studies.map((study, index) => {
     const h: HistoryPlotInfo = {
       study_name: study.name,
@@ -64,8 +70,8 @@ export const PlotHistory: FC<{
     return h
   })
 
-  const handleObjectiveChange = (event: SelectChangeEvent<number>) => {
-    setObjectiveId(event.target.value as number)
+  const handleObjectiveChange = (event: SelectChangeEvent<string>) => {
+    setTarget(event.target.value)
   }
 
   const handleXAxisChange = (e: ChangeEvent<HTMLInputElement>) => {
@@ -78,33 +84,18 @@ export const PlotHistory: FC<{
     }
   }
 
-  const handleLogScaleChange = () => {
-    setLogScale(!logScale)
-  }
-
-  const handleFilterPrunedChange = () => {
-    setFilterPrunedTrial(!filterPrunedTrial)
-  }
-
   useEffect(() => {
     if (graphComponentState !== "componentWillMount") {
       plotHistory(
         historyPlotInfos,
-        target,
+        selected,
         xAxis,
         logScale,
-        theme.palette.mode,
-        colorTheme,
-
+        colorThemeUsed,
+        markerSize
       )
     }
-  }, [
-    studies,
-    target,
-    logScale,
-    xAxis,
-    theme.palette.mode,
-  ])
+  }, [graphComponentState, studies, selected, logScale, xAxis, colorThemeUsed, markerSize])
 
   return (
     <Grid container direction="row">
@@ -115,60 +106,30 @@ export const PlotHistory: FC<{
         direction="column"
         sx={{ paddingRight: theme.spacing(2) }}
       >
-        <Typography variant="h6" sx={{ margin: "1em 0", fontWeight: 600 }}>
+        <Typography
+          variant="h6"
+          sx={{ margin: "1em 0", fontWeight: theme.typography.fontWeightBold }}
+        >
           History
         </Typography>
-        {study !== null && study.directions.length !== 1 ? (
+        {studies[0] !== null && targets.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) => (
-                // biome-ignore lint/suspicious/noArrayIndexKey: <explanation>
-                <MenuItem value={i} key={i}>
-                  {i}
+            <FormLabel component="legend">y Axis</FormLabel>
+            <Select
+              value={selected.identifier()}
+              onChange={handleObjectiveChange}
+            >
+              {targets.map((t, i) => (
+                <MenuItem value={t.identifier()} key={i}>
+                  {t.toLabel(studies[0].metric_names)}
                 </MenuItem>
               ))}
             </Select>
           </FormControl>
         ) : null}
-        <FormControl
-          component="fieldset"
-          sx={{ marginBottom: theme.spacing(2) }}
-        >
-          <FormLabel component="legend">Log y scale:</FormLabel>
-          <Switch
-            checked={logScale}
-            onChange={handleLogScaleChange}
-            value="enable"
-          />
-        </FormControl>
-        <FormControl
-          component="fieldset"
-          sx={{ marginBottom: theme.spacing(2) }}
-        >
-          <FormLabel component="legend">Filter state:</FormLabel>
-          <FormControlLabel
-            control={
-              <Checkbox
-                checked={!filterCompleteTrial}
-                onChange={handleFilterCompleteChange}
-              />
-            }
-            label="Complete"
-          />
-          <FormControlLabel
-            control={
-              <Checkbox
-                checked={!filterPrunedTrial}
-                onChange={handleFilterPrunedChange}
-              />
-            }
-            label="Pruned"
-          />
-        </FormControl>
         <FormControl
           component="fieldset"
           sx={{ marginBottom: theme.spacing(2) }}
@@ -197,6 +158,20 @@ export const PlotHistory: FC<{
             />
           </RadioGroup>
         </FormControl>
+        <FormControl>
+          <FormLabel component="legend">Marker size:</FormLabel>
+          <Slider
+            defaultValue={5}
+            marks={true}
+            min={1}
+            max={10}
+            step={1}
+            onChange={(e) => {
+              // @ts-ignore
+              setMarkerSize(e.target.value as number)
+            }}
+          />
+        </FormControl>
       </Grid>
       <Grid item xs={9}>
         <div id={plotDomId} />
@@ -210,7 +185,6 @@ const plotHistory = (
   target: Target,
   xAxis: "number" | "datetime_start" | "datetime_complete",
   logScale: boolean,
-  mode: string,
   colorTheme: Partial<Plotly.Template>,
   markerSize: number
 ) => {
@@ -232,7 +206,7 @@ const plotHistory = (
       b: 0,
     },
     yaxis: {
-      title: target.toLabel(historyPlotInfos[0].objective_names),
+      title: target.toLabel(historyPlotInfos[0].metric_names),
       type: logScale ? "log" : "linear",
     },
     xaxis: {
@@ -257,22 +231,22 @@ const plotHistory = (
 
   const plotData: Partial<plotly.PlotData>[] = []
   const infeasiblePlotData: Partial<plotly.PlotData>[] = []
-  historyPlotInfos.forEach((h) => {
+  for (const h of historyPlotInfos) {
     const feasibleTrials: Optuna.Trial[] = []
     const infeasibleTrials: Optuna.Trial[] = []
-    h.trials.forEach((t) => {
+    for (const t of h.trials) {
       if (t.constraints.every((c) => c <= 0)) {
         feasibleTrials.push(t)
       } else {
         infeasibleTrials.push(t)
       }
-    })
+    }
     plotData.push({
       x: feasibleTrials.map(getAxisX),
       y: feasibleTrials.map(
         (t: Optuna.Trial): number => target.getTargetValue(t) as number
       ),
-      name: `${target.toLabel(h.objective_names)} of ${h.study_name}`,
+      name: `${target.toLabel(h.metric_names)} of ${h.study_name}`,
       marker: {
         size: markerSize,
       },
@@ -290,7 +264,9 @@ const plotHistory = (
         const value = target.getTargetValue(t) as number
         if (t.state !== "Complete") {
           continue
-        } else if (currentBest === null) {
+        }
+
+        if (currentBest === null) {
           currentBest = value
           xForLinePlot.push(getAxisX(t))
           yForLinePlot.push(value)
@@ -340,13 +316,13 @@ const plotHistory = (
       name: `Infeasible Trial of ${h.study_name}`,
       marker: {
         size: markerSize,
-        color: mode === "dark" ? "#666666" : "#cccccc",
+        color: colorTheme === plotlyDarkTemplate ? "#666666" : "#cccccc",
       },
       mode: "markers",
       type: "scatter",
       showlegend: false,
     })
-  })
+  }
   plotData.push(...infeasiblePlotData)
   plotly.react(plotDomId, plotData, layout)
 }