Refactor hyperparameter importance

optuna_dashboard/_app.py

@@ -336,20 +336,19 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle:
     @app.get("/api/studies/<study_id:int>/param_importances")
     @json_api_view
     def get_param_importances(study_id: int) -> BottleViewReturn:
-        # TODO(chenghuzi): add support for selecting params via query parameters.
-        objective_id = int(request.params.get("objective_id", 0))
         try:
             n_directions = len(storage.get_study_directions(study_id))
         except KeyError:
             response.status = 404  # Study is not found
             return {"reason": f"study_id={study_id} is not found"}
-        if objective_id >= n_directions:
-            response.status = 400  # Bad request
-            return {"reason": f"study_id={study_id} has only {n_directions} direction(s)."}
 
         trials = get_trials(storage, study_id)
         try:
-            return get_param_importance_from_trials_cache(storage, study_id, objective_id, trials)
+            importances = [
+                get_param_importance_from_trials_cache(storage, study_id, objective_id, trials)
+                for objective_id in range(n_directions)
+            ]
+            return {"param_importances": importances}
         except ValueError as e:
             response.status = 400  # Bad request
             return {"reason": str(e)}

optuna_dashboard/_importance.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import threading
+from typing import List
 from typing import TYPE_CHECKING
 import warnings
 
@@ -25,26 +26,18 @@ except Exception as e:
 if TYPE_CHECKING:
     from typing import TypedDict
 
-    ImportanceItemType = TypedDict(
-        "ImportanceItemType",
+    ImportanceType = TypedDict(
+        "ImportanceType",
         {
             "name": str,
             "importance": float,
             "distribution": str,
         },
     )
-    ImportanceType = TypedDict(
-        "ImportanceType",
-        {
-            "target_name": str,
-            "param_importances": list[ImportanceItemType],
-        },
-    )
 
-target_name = "Objective Value"
 param_importance_cache_lock = threading.Lock()
 # { "{study_id}:{objective_id}" : (n_completed_trials, importance) }
-param_importance_cache: dict[str, tuple[int, ImportanceType]] = {}
+param_importance_cache: dict[str, tuple[int, list[ImportanceType]]] = {}
 
 
 class StudyWrapper(Study):
@@ -62,16 +55,16 @@ class StudyWrapper(Study):
 
 def get_param_importance_from_trials_cache(
     storage: BaseStorage, study_id: int, objective_id: int, trials: list[FrozenTrial]
-) -> ImportanceType:
+) -> list[ImportanceType]:
     completed_trials = [t for t in trials if t.state == TrialState.COMPLETE]
     n_completed_trials = len(completed_trials)
     if n_completed_trials == 0:
-        return {"target_name": target_name, "param_importances": []}
+        return []
 
     cache_key = f"{study_id}:{objective_id}"
     with param_importance_cache_lock:
         cache_n_trial, cache_importance = param_importance_cache.get(
-            cache_key, (0, {"target_name": target_name, "param_importances": []})
+            cache_key, (0, {"param_importances": []})
         )
         if n_completed_trials == cache_n_trial:
             return cache_importance
@@ -95,18 +88,15 @@ def get_param_importance_from_trials_cache(
 
 def convert_to_importance_type(
     importance: dict[str, float], trials: list[FrozenTrial]
-) -> ImportanceType:
-    return {
-        "target_name": target_name,
-        "param_importances": [
-            {
-                "name": name,
-                "importance": importance,
-                "distribution": get_distribution_name(name, trials),
-            }
-            for name, importance in importance.items()
-        ],
-    }
+) -> list[ImportanceType]:
+    return [
+        {
+            "name": name,
+            "importance": importance,
+            "distribution": get_distribution_name(name, trials),
+        }
+        for name, importance in importance.items()
+    ]
 
 
 def get_distribution_name(param_name: str, trials: list[FrozenTrial]) -> str:

optuna_dashboard/ts/action.ts

@@ -3,6 +3,7 @@ import { useSnackbar } from "notistack"
 import {
   getStudyDetailAPI,
   getStudySummariesAPI,
+  getParamImportances,
   createNewStudyAPI,
   deleteStudyAPI,
   saveNoteAPI,
@@ -11,6 +12,7 @@ import {
   graphVisibilityState,
   studyDetailsState,
   studySummariesState,
+  paramImportanceState,
 } from "./state"
 
 const localStorageGraphVisibility = "graphVisibility"
@@ -23,6 +25,8 @@ export const actionCreator = () => {
     useRecoilState<StudyDetails>(studyDetailsState)
   const [graphVisibility, setGraphVisibility] =
     useRecoilState<GraphVisibility>(graphVisibilityState)
+  const [paramImportance, setParamImportance] =
+    useRecoilState<StudyParamImportance>(paramImportanceState)
 
   const setStudyDetailState = (studyId: number, study: StudyDetail) => {
     const newVal = Object.assign({}, studyDetails)
@@ -30,6 +34,15 @@ export const actionCreator = () => {
     setStudyDetails(newVal)
   }
 
+  const setStudyParamImportanceState = (
+    studyId: number,
+    importance: ParamImportance[][]
+  ) => {
+    const newVal = Object.assign({}, paramImportance)
+    newVal[studyId] = importance
+    setParamImportance(newVal)
+  }
+
   const updateStudySummaries = (successMsg?: string) => {
     getStudySummariesAPI()
       .then((studySummaries: StudySummary[]) => {
@@ -77,6 +90,22 @@ export const actionCreator = () => {
       })
   }
 
+  const updateParamImportance = (studyId: number) => {
+    getParamImportances(studyId)
+      .then((importance) => {
+        setStudyParamImportanceState(studyId, importance)
+      })
+      .catch((err) => {
+        const reason = err.response?.data.reason
+        enqueueSnackbar(
+          `Failed to load hyperparameter importance (reason=${reason})`,
+          {
+            variant: "error",
+          }
+        )
+      })
+  }
+
   const createNewStudy = (studyName: string, directions: StudyDirection[]) => {
     createNewStudyAPI(studyName, directions)
       .then((study_summary) => {
@@ -157,6 +186,7 @@ export const actionCreator = () => {
   return {
     updateStudyDetail,
     updateStudySummaries,
+    updateParamImportance,
     createNewStudy,
     deleteStudy,
     getGraphVisibility,

optuna_dashboard/ts/apiClient.ts

@@ -195,24 +195,15 @@ export const saveNoteAPI = (
 }
 
 interface ParamImportancesResponse {
-  target_name: string
-  param_importances: ParamImportance[]
+  param_importances: ParamImportance[][]
 }
 
 export const getParamImportances = (
-  studyId: number,
-  objectiveId = 0
-): Promise<ParamImportances> => {
+  studyId: number
+): Promise<ParamImportance[][]> => {
   return axiosInstance
-    .get<ParamImportancesResponse>(
-      `/api/studies/${studyId}/param_importances`,
-      {
-        params: {
-          objective_id: objectiveId,
-        },
-      }
-    )
+    .get<ParamImportancesResponse>(`/api/studies/${studyId}/param_importances`)
     .then((res) => {
-      return res.data
+      return res.data.param_importances
     })
 }

optuna_dashboard/ts/components/GraphHyperparameterImportances.tsx

@@ -1,5 +1,6 @@
 import * as plotly from "plotly.js-dist-min"
 import React, { FC, useEffect, useState } from "react"
+import { useRecoilValue } from "recoil"
 import {
   Grid,
   FormControl,
@@ -12,49 +13,43 @@ import {
   Box,
 } from "@mui/material"
 
-import { getParamImportances } from "../apiClient"
 import { plotlyDarkTemplate } from "./PlotlyDarkMode"
-import { useSnackbar } from "notistack"
+import { actionCreator } from "../action"
+import { paramImportanceState } from "../state"
 const plotDomId = "graph-hyperparameter-importances"
 
+const useParamImportanceValue = (
+  studyId: number
+): ParamImportance[][] | null => {
+  const studyParamImportance =
+    useRecoilValue<StudyParamImportance>(paramImportanceState)
+  return studyParamImportance[studyId] || null
+}
+
 export const GraphHyperparameterImportances: FC<{
   study: StudyDetail | null
   studyId: number
 }> = ({ study = null, studyId }) => {
   const theme = useTheme()
+  const action = actionCreator()
+  const importances = useParamImportanceValue(studyId)
   const [objectiveId, setObjectiveId] = useState<number>(0)
   const numCompletedTrials =
     study?.trials.filter((t) => t.state === "Complete").length || 0
-  const [importances, setImportances] = useState<ParamImportances | null>(null)
-  const { enqueueSnackbar } = useSnackbar()
 
   const handleObjectiveChange = (event: SelectChangeEvent<number>) => {
     setObjectiveId(event.target.value as number)
   }
 
   useEffect(() => {
-    if (numCompletedTrials > 0) {
-      getParamImportances(studyId, objectiveId)
-        .then((p) => {
-          setImportances(p)
-        })
-        .catch((err) => {
-          const reason = err.response?.data.reason
-          enqueueSnackbar(
-            `Failed to load hyperparameter importance (reason=${reason})`,
-            {
-              variant: "error",
-            }
-          )
-        })
-    }
-  }, [numCompletedTrials, objectiveId, theme.palette.mode])
+    action.updateParamImportance(studyId)
+  }, [numCompletedTrials])
 
   useEffect(() => {
-    if (importances !== null) {
-      plotParamImportances(importances, theme.palette.mode)
+    if (importances !== null && importances.length > objectiveId) {
+      plotParamImportances(importances[objectiveId], theme.palette.mode)
     }
-  }, [importances, theme.palette.mode])
+  }, [importances, objectiveId, theme.palette.mode])
 
   return (
     <Grid container direction="row">
@@ -88,25 +83,20 @@ export const GraphHyperparameterImportances: FC<{
   )
 }
 
-const plotParamImportances = (
-  paramsImportanceData: ParamImportances,
-  mode: string
-) => {
+const plotParamImportances = (importance: ParamImportance[], mode: string) => {
   if (document.getElementById(plotDomId) === null) {
     return
   }
-  const param_importances = [
-    ...paramsImportanceData.param_importances,
-  ].reverse()
-  const importance_values = param_importances.map((p) => p.importance)
-  const param_names = param_importances.map((p) => p.name)
-  const param_hover_templates = param_importances.map(
+  const reversed = [...importance].reverse()
+  const importance_values = reversed.map((p) => p.importance)
+  const param_names = reversed.map((p) => p.name)
+  const param_hover_templates = reversed.map(
     (p) => `${p.name} (${p.distribution}): ${p.importance} <extra></extra>`
   )
 
   const layout: Partial<plotly.Layout> = {
     xaxis: {
-      title: `Importance for ${paramsImportanceData.target_name}`,
+      title: `Importance for the Objective Value`,
     },
     yaxis: {
       title: "Hyperparameter",

optuna_dashboard/ts/state.ts

@@ -10,6 +10,11 @@ export const studyDetailsState = atom<StudyDetails>({
   default: {},
 })
 
+export const paramImportanceState = atom<StudyParamImportance>({
+  key: "paramImportance",
+  default: {},
+})
+
 export const graphVisibilityState = atom<GraphVisibility>({
   key: "graphVisibility",
   default: {

optuna_dashboard/ts/types/index.d.ts

@@ -109,7 +109,6 @@ declare interface StudyDetails {
   [study_id: string]: StudyDetail
 }
 
-declare interface ParamImportances {
-  target_name: string
-  param_importances: ParamImportance[]
+declare interface StudyParamImportance {
+  [study_id: string]: ParamImportance[][]
 }