Add includePruned as optional

optuna · Aug 20, 2024 · 6498a55 · 6498a55
1 parent 947c5bb
commit 6498a55
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git a/optuna_dashboard/ts/components/GraphHistory.tsx b/optuna_dashboard/ts/components/GraphHistory.tsx
@@ -8,13 +8,14 @@ export const GraphHistory: FC<{
   studies: StudyDetail[]
   logScale: boolean
   includePruned: boolean
-}> = ({ studies, logScale }) => {
+}> = ({ studies, logScale, includePruned }) => {
   const theme = useTheme()
   const colorTheme = usePlotlyColorTheme(theme.palette.mode)
   return (
     <PlotHistory
       studies={studies}
       logScale={logScale}
+      includePruned={includePruned}
       colorTheme={colorTheme}
     />
   )

diff --git a/tslib/react/src/components/PlotHistory.tsx b/tslib/react/src/components/PlotHistory.tsx
@@ -37,10 +37,9 @@ interface HistoryPlotInfo {
 export const PlotHistory: FC<{
   studies: Optuna.Study[]
   logScale?: boolean
+  includePruned?: boolean
   colorTheme?: Partial<Plotly.Template>
-}> = ({ studies, logScale, colorTheme }) => {
-  const filterPrunedTrial = false
-
+}> = ({ studies, logScale, includePruned, colorTheme }) => {
   const { graphComponentState, notifyGraphDidRender } = useGraphComponentState()
 
   const theme = useTheme()
@@ -52,6 +51,8 @@ export const PlotHistory: FC<{
   >("number")
 
   const [logScaleInternal, setLogScaleInternal] = useState<boolean>(false)
+  const [includePrunedInternal, setIncludePrunedInternal] =
+    useState<boolean>(true)
 
   const [markerSize, setMarkerSize] = useState<number>(5)
 
@@ -61,7 +62,7 @@ export const PlotHistory: FC<{
   const trials = useFilteredTrialsFromStudies(
     studies,
     [selected],
-    filterPrunedTrial
+    includePruned === undefined ? !includePrunedInternal : !includePruned
   )
 
   const historyPlotInfos = studies.map((study, index) => {
@@ -82,6 +83,10 @@ export const PlotHistory: FC<{
     setLogScaleInternal(!logScaleInternal)
   }
 
+  const handleIncludePrunedChange = () => {
+    setIncludePrunedInternal(!includePrunedInternal)
+  }
+
   const handleXAxisChange = (e: ChangeEvent<HTMLInputElement>) => {
     if (e.target.value === "number") {
       setXAxis("number")
@@ -189,6 +194,19 @@ export const PlotHistory: FC<{
             />
           </FormControl>
         ) : null}
+        {includePruned === undefined ? (
+          <FormControl
+            component="fieldset"
+            sx={{ marginBottom: theme.spacing(2) }}
+          >
+            <FormLabel component="legend"> Include PRUNED trials:</FormLabel>
+            <Switch
+              checked={includePrunedInternal}
+              onChange={handleIncludePrunedChange}
+              value="enable"
+            />
+          </FormControl>
+        ) : null}
         <FormControl>
           <FormLabel component="legend">Marker size:</FormLabel>
           <Slider