diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 8183f48d..d10bc9c1 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -336,20 +336,19 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle: @app.get("/api/studies//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)} diff --git a/optuna_dashboard/_importance.py b/optuna_dashboard/_importance.py index f6e0e4fa..ec6507e1 100644 --- a/optuna_dashboard/_importance.py +++ b/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: diff --git a/optuna_dashboard/ts/action.ts b/optuna_dashboard/ts/action.ts index da87dc3f..335a815c 100644 --- a/optuna_dashboard/ts/action.ts +++ b/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(studyDetailsState) const [graphVisibility, setGraphVisibility] = useRecoilState(graphVisibilityState) + const [paramImportance, setParamImportance] = + useRecoilState(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, diff --git a/optuna_dashboard/ts/apiClient.ts b/optuna_dashboard/ts/apiClient.ts index d7812a56..405e18f6 100644 --- a/optuna_dashboard/ts/apiClient.ts +++ b/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 => { + studyId: number +): Promise => { return axiosInstance - .get( - `/api/studies/${studyId}/param_importances`, - { - params: { - objective_id: objectiveId, - }, - } - ) + .get(`/api/studies/${studyId}/param_importances`) .then((res) => { - return res.data + return res.data.param_importances }) } diff --git a/optuna_dashboard/ts/components/GraphHyperparameterImportances.tsx b/optuna_dashboard/ts/components/GraphHyperparameterImportances.tsx index fa7451e3..9cfa08b1 100644 --- a/optuna_dashboard/ts/components/GraphHyperparameterImportances.tsx +++ b/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(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(0) const numCompletedTrials = study?.trials.filter((t) => t.state === "Complete").length || 0 - const [importances, setImportances] = useState(null) - const { enqueueSnackbar } = useSnackbar() const handleObjectiveChange = (event: SelectChangeEvent) => { 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 ( @@ -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} ` ) const layout: Partial = { xaxis: { - title: `Importance for ${paramsImportanceData.target_name}`, + title: `Importance for the Objective Value`, }, yaxis: { title: "Hyperparameter", diff --git a/optuna_dashboard/ts/state.ts b/optuna_dashboard/ts/state.ts index e0ea82bb..5dda1de3 100644 --- a/optuna_dashboard/ts/state.ts +++ b/optuna_dashboard/ts/state.ts @@ -10,6 +10,11 @@ export const studyDetailsState = atom({ default: {}, }) +export const paramImportanceState = atom({ + key: "paramImportance", + default: {}, +}) + export const graphVisibilityState = atom({ key: "graphVisibility", default: { diff --git a/optuna_dashboard/ts/types/index.d.ts b/optuna_dashboard/ts/types/index.d.ts index 1cb85277..c5a64b57 100644 --- a/optuna_dashboard/ts/types/index.d.ts +++ b/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[][] }