Move PlotTimeline from optuna_dashboard/ts to tslib/react

optuna_dashboard/ts/components/GraphTimeline.tsx

@@ -1,24 +1,28 @@
-import { Card, CardContent, Grid, Typography, useTheme } from "@mui/material"
-import * as Optuna from "@optuna/types"
+import { Grid, useTheme } from "@mui/material"
+import { PlotTimeline } from "@optuna/react"
 import * as plotly from "plotly.js-dist-min"
 import React, { FC, useEffect } from "react"
-import { StudyDetail, Trial } from "ts/types/optuna"
+import { StudyDetail } from "ts/types/optuna"
 import { PlotType } from "../apiClient"
-import { makeHovertext } from "../graphUtil"
+import { studyDetailToStudy } from "../graphUtil"
 import { usePlot } from "../hooks/usePlot"
 import { usePlotlyColorTheme } from "../state"
 import { useBackendRender } from "../state"
 
 const plotDomId = "graph-timeline"
-const maxBars = 100
 
 export const GraphTimeline: FC<{
   study: StudyDetail | null
 }> = ({ study }) => {
+  const theme = useTheme()
+  const colorTheme = usePlotlyColorTheme(theme.palette.mode)
+
   if (useBackendRender()) {
     return <GraphTimelineBackend study={study} />
   } else {
-    return <GraphTimelineFrontend study={study} />
+    return (
+      <PlotTimeline study={studyDetailToStudy(study)} colorTheme={colorTheme} />
+    )
   }
 }
 
@@ -51,176 +55,3 @@ const GraphTimelineBackend: FC<{
     </Grid>
   )
 }
-
-const GraphTimelineFrontend: FC<{
-  study: StudyDetail | null
-}> = ({ study }) => {
-  const theme = useTheme()
-  const colorTheme = usePlotlyColorTheme(theme.palette.mode)
-
-  const trials = study?.trials ?? []
-
-  useEffect(() => {
-    if (study !== null) {
-      plotTimeline(trials, colorTheme)
-    }
-  }, [trials, colorTheme])
-
-  return (
-    <Card>
-      <CardContent>
-        <Typography
-          variant="h6"
-          sx={{ margin: "1em 0", fontWeight: theme.typography.fontWeightBold }}
-        >
-          Timeline
-        </Typography>
-        <Grid item xs={9}>
-          <div id={plotDomId} />
-        </Grid>
-      </CardContent>
-    </Card>
-  )
-}
-
-const plotTimeline = (
-  trials: Trial[],
-  colorTheme: Partial<Plotly.Template>
-) => {
-  if (document.getElementById(plotDomId) === null) {
-    return
-  }
-
-  if (trials.length === 0) {
-    plotly.react(plotDomId, [], {
-      template: colorTheme,
-    })
-    return
-  }
-
-  const cm: Record<Optuna.TrialState, string> = {
-    Complete: "blue",
-    Fail: "red",
-    Pruned: "orange",
-    Running: "green",
-    Waiting: "gray",
-  }
-  const runningKey = "Running"
-
-  const lastTrials = trials.slice(-maxBars) // To only show last elements
-  const minDatetime = new Date(
-    Math.min(
-      ...lastTrials.map(
-        (t) => t.datetime_start?.getTime() ?? new Date().getTime()
-      )
-    )
-  )
-  const maxRunDuration = Math.max(
-    ...trials.map((t) => {
-      return t.datetime_start === undefined || t.datetime_complete === undefined
-        ? -Infinity
-        : t.datetime_complete.getTime() - t.datetime_start.getTime()
-    })
-  )
-  const hasRunning =
-    (maxRunDuration === -Infinity &&
-      trials.some((t) => t.state === runningKey)) ||
-    trials.some((t) => {
-      if (t.state !== runningKey) {
-        return false
-      }
-      const now = new Date().getTime()
-      const start = t.datetime_start?.getTime() ?? now
-      // This is an ad-hoc handling to check if the trial is running.
-      // We do not check via `trialState` because some trials may have state=RUNNING,
-      // even if they are not running because of unexpected job kills.
-      // In this case, we would like to ensure that these trials will not squash the timeline plot
-      // for the other trials.
-      return now - start < maxRunDuration * 5
-    })
-  const maxDatetime = hasRunning
-    ? new Date()
-    : new Date(
-        Math.max(
-          ...lastTrials.map(
-            (t) => t.datetime_complete?.getTime() ?? minDatetime.getTime()
-          )
-        )
-      )
-  const layout: Partial<plotly.Layout> = {
-    margin: {
-      l: 50,
-      t: 0,
-      r: 50,
-      b: 0,
-    },
-    xaxis: {
-      title: "Datetime",
-      type: "date",
-      range: [minDatetime, maxDatetime],
-    },
-    yaxis: {
-      title: "Trial",
-      range: [lastTrials[0].number, lastTrials[0].number + lastTrials.length],
-    },
-    uirevision: "true",
-    template: colorTheme,
-    legend: {
-      x: 1.0,
-      y: 0.95,
-    },
-  }
-
-  const makeTrace = (bars: Trial[], state: string, color: string) => {
-    const isRunning = state === runningKey
-    // Waiting trials should not squash other trials, so use `maxDatetime` instead of `new Date()`.
-    const starts = bars.map((b) => b.datetime_start ?? maxDatetime)
-    const runDurations = bars.map((b, i) => {
-      const startTime = starts[i].getTime()
-      const completeTime = isRunning
-        ? maxDatetime.getTime()
-        : b.datetime_complete?.getTime() ?? startTime
-      // By using 1 as the min value, we can recognize these bars at least when zooming in.
-      return Math.max(1, completeTime - startTime)
-    })
-    const trace: Partial<plotly.PlotData> = {
-      type: "bar",
-      x: runDurations,
-      y: bars.map((b) => b.number),
-      // @ts-ignore: To suppress ts(2322)
-      base: starts,
-      name: state,
-      text: bars.map((b) => makeHovertext(b)),
-      hovertemplate: "%{text}<extra>" + state + "</extra>",
-      orientation: "h",
-      marker: { color: color },
-      textposition: "none", // Avoid drawing hovertext in a bar.
-    }
-    return trace
-  }
-
-  const traces: Partial<plotly.PlotData>[] = []
-  for (const [state, color] of Object.entries(cm)) {
-    const bars = trials.filter((t) => t.state === state)
-    if (bars.length === 0) {
-      continue
-    }
-    if (state === "Complete") {
-      const feasibleTrials = bars.filter((t) =>
-        t.constraints.every((c) => c <= 0)
-      )
-      const infeasibleTrials = bars.filter((t) =>
-        t.constraints.some((c) => c > 0)
-      )
-      if (feasibleTrials.length > 0) {
-        traces.push(makeTrace(feasibleTrials, "Complete", color))
-      }
-      if (infeasibleTrials.length > 0) {
-        traces.push(makeTrace(infeasibleTrials, "Infeasible", "#cccccc"))
-      }
-    } else {
-      traces.push(makeTrace(bars, state, color))
-    }
-  }
-  plotly.react(plotDomId, traces, layout)
-}

optuna_dashboard/ts/components/StudyDetail.tsx

@@ -16,9 +16,9 @@ import { Link, useParams } from "react-router-dom"
 import { useRecoilValue } from "recoil"
 
 import { TrialTable } from "@optuna/react"
-import * as Optuna from "@optuna/types"
 import { actionCreator } from "../action"
 import { useConstants } from "../constantsProvider"
+import { studyDetailToStudy } from "../graphUtil"
 import {
   reloadIntervalState,
   useStudyDetailValue,
@@ -62,19 +62,7 @@ export const StudyDetail: FC<{
   const reloadInterval = useRecoilValue<number>(reloadIntervalState)
   const studyName = useStudyName(studyId)
   const isPreferential = useStudyIsPreferential(studyId)
-  const study: Optuna.Study | null = studyDetail
-    ? {
-        id: studyDetail.id,
-        name: studyDetail.name,
-        directions: studyDetail.directions,
-        union_search_space: studyDetail.union_search_space,
-        intersection_search_space: studyDetail.intersection_search_space,
-        union_user_attrs: studyDetail.union_user_attrs,
-        datetime_start: studyDetail.datetime_start,
-        trials: studyDetail.trials,
-        metric_names: studyDetail.metric_names,
-      }
-    : null
+  const study = studyDetailToStudy(studyDetail)
 
   const title =
     studyName !== null ? `${studyName} (id=${studyId})` : `Study #${studyId}`

optuna_dashboard/ts/graphUtil.ts

@@ -1,5 +1,5 @@
 import * as Optuna from "@optuna/types"
-import { SearchSpaceItem, Trial } from "./types/optuna"
+import { SearchSpaceItem, StudyDetail, Trial } from "./types/optuna"
 
 const PADDING_RATIO = 0.05
 
@@ -115,3 +115,23 @@ export const makeHovertext = (trial: Trial): string => {
     "  "
   ).replace(/\n/g, "<br>")
 }
+
+export const studyDetailToStudy = (
+  studyDetail: StudyDetail | null
+): Optuna.Study | null => {
+  const study: Optuna.Study | null = studyDetail
+    ? {
+        id: studyDetail.id,
+        name: studyDetail.name,
+        directions: studyDetail.directions,
+        union_search_space: studyDetail.union_search_space,
+        intersection_search_space: studyDetail.intersection_search_space,
+        union_user_attrs: studyDetail.union_user_attrs,
+        datetime_start: studyDetail.datetime_start,
+        trials: studyDetail.trials,
+        metric_names: studyDetail.metric_names,
+      }
+    : null
+
+  return study
+}

tslib/react/src/components/PlotTimeline.stories.tsx

@@ -0,0 +1,53 @@
+import { CssBaseline, ThemeProvider } from "@mui/material"
+import { Meta, StoryObj } from "@storybook/react"
+import React from "react"
+import { useMockStudy } from "../MockStudies"
+import { darkTheme } from "../styles/darkTheme"
+import { lightTheme } from "../styles/lightTheme"
+import { PlotTimeline } from "./PlotTimeline"
+
+const meta: Meta<typeof PlotTimeline> = {
+  component: PlotTimeline,
+  title: "Plot/Timeline",
+  tags: ["autodocs"],
+  decorators: [
+    (Story, storyContext) => {
+      const { study } = useMockStudy(storyContext.parameters?.studyId)
+      if (!study) return <p>loading...</p>
+      return (
+        <ThemeProvider theme={storyContext.parameters?.theme}>
+          <CssBaseline />
+          <Story
+            args={{
+              study,
+            }}
+          />
+        </ThemeProvider>
+      )
+    },
+  ],
+}
+
+export default meta
+type Story = StoryObj<typeof PlotTimeline>
+
+export const LightTheme: Story = {
+  parameters: {
+    studyId: 1,
+    theme: lightTheme,
+  },
+}
+
+export const DarkTheme: Story = {
+  parameters: {
+    studyId: 1,
+    theme: darkTheme,
+  },
+}
+
+// TODO(c-bata): Add a story for multi objective study.
+// export const MultiObjective: Story = {
+//   parameters: {
+//     ...
+//   },
+// }