From e39357bd20a89d804507019a091b3fdc077cf057 Mon Sep 17 00:00:00 2001 From: c-bata Date: Wed, 6 Sep 2023 18:01:32 +0900 Subject: [PATCH] Support user-defined plotly figures --- optuna_dashboard/__init__.py | 1 + optuna_dashboard/_app.py | 4 + optuna_dashboard/_custom_plot_data.py | 115 ++++++++++++++++++ optuna_dashboard/_serializer.py | 5 + optuna_dashboard/ts/apiClient.ts | 2 + .../ts/components/StudyHistory.tsx | 12 ++ .../ts/components/UserDefinedPlot.tsx | 16 +++ optuna_dashboard/ts/types/index.d.ts | 6 + pyproject.toml | 1 + python_tests/test_custom_plot_data.py | 63 ++++++++++ python_tests/test_serializers.py | 4 +- 11 files changed, 227 insertions(+), 2 deletions(-) create mode 100644 optuna_dashboard/_custom_plot_data.py create mode 100644 optuna_dashboard/ts/components/UserDefinedPlot.tsx create mode 100644 python_tests/test_custom_plot_data.py diff --git a/optuna_dashboard/__init__.py b/optuna_dashboard/__init__.py index 3d363cf4..493af736 100644 --- a/optuna_dashboard/__init__.py +++ b/optuna_dashboard/__init__.py @@ -1,5 +1,6 @@ from ._app import run_server # noqa from ._app import wsgi # noqa +from ._custom_plot_data import save_plotly_graph_object # noqa from ._form_widget import ChoiceWidget # noqa from ._form_widget import dict_to_form_widget # noqa from ._form_widget import ObjectiveChoiceWidget # noqa diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 5f433072..c32c2061 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -25,6 +25,7 @@ from . import _note as note from ._bottle_util import BottleViewReturn from ._bottle_util import json_api_view from ._cached_extra_study_property import get_cached_extra_study_property +from ._custom_plot_data import get_plotly_graph_objects from ._importance import get_param_importance_from_trials_cache from ._pareto_front import get_pareto_front_trials from ._preferential_history import NewHistory @@ -214,6 +215,8 @@ def create_app( union_user_attrs, has_intermediate_values, ) = get_cached_extra_study_property(study_id, trials) + + plotly_graph_objects = get_plotly_graph_objects(system_attrs) return serialize_study_detail( summary, best_trials, @@ -222,6 +225,7 @@ def create_app( union, union_user_attrs, has_intermediate_values, + plotly_graph_objects, ) @app.get("/api/studies//param_importances") diff --git a/optuna_dashboard/_custom_plot_data.py b/optuna_dashboard/_custom_plot_data.py new file mode 100644 index 00000000..d669d4ad --- /dev/null +++ b/optuna_dashboard/_custom_plot_data.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +import math +from typing import TYPE_CHECKING +import uuid + +from optuna import Study + + +if TYPE_CHECKING: + from typing import Any + + from optuna.storages import BaseStorage + import plotly.graph_objs as go + + +SYSTEM_ATTR_PLOT_DATA = "dashboard:plot_data:" +SYSTEM_ATTR_MAX_LENGTH = 2045 + + +def save_plotly_graph_object( + study: Study, figure: go.Figure, *, graph_object_id: str | None = None +) -> str: + """Save the user-defined plotly's graph object to the study. + + Example: + + .. code-block:: python + + import optuna + from optuna_dashboard import save_plotly_graph_object + + def objective(trial): + x = trial.suggest_float("x", -100, 100) + y = trial.suggest_categorical("y", [-1, 0, 1]) + return x**2 + y + + study = optuna.create_study() + study.optimize(objective, n_trials=100) + + figure = optuna.visualization.plot_optimization_history(study) + save_plotly_graph_object(study, figure) + + Args: + study: + Target study object. + plot_data: + The plotly's graph object to save. + graph_object_id: + Unique identifier of the graph object. If specified, the graph object is overwritten. + + Returns: + The graph object ID. + """ + storage = study._storage + study_id = study._study_id + + graph_object_id = graph_object_id or str(uuid.uuid4()) + key = SYSTEM_ATTR_PLOT_DATA + graph_object_id + ":" + plot_data_json_str = figure.to_json() + save_graph_object_json(storage, study_id, key, plot_data_json_str) + return graph_object_id + + +def save_graph_object_json( + storage: BaseStorage, study_id: int, key_prefix: str, plot_data_json_str: str +) -> None: + plot_data_system_attrs = split_plot_data(plot_data_json_str, key_prefix) + for k, v in plot_data_system_attrs.items(): + storage.set_study_system_attr(study_id, k, v) + + # Clear previous graph object attributes + study_system_attrs = storage.get_study_system_attrs(study_id) + all_plot_data_system_attrs = [k for k in study_system_attrs if k.startswith(key_prefix)] + if len(all_plot_data_system_attrs) > len(plot_data_system_attrs): + for i in range(len(plot_data_system_attrs), len(all_plot_data_system_attrs)): + storage.set_study_system_attr(study_id, f"{key_prefix}{i}", "") + + +def list_graph_object_ids(system_attrs: dict[str, Any]) -> list[str]: + titles = set() + for key in system_attrs: + if not key.startswith(SYSTEM_ATTR_PLOT_DATA): + continue + + s = key.split(":", maxsplit=2) # e.g. ["dashboard", "plot_data", "Optimization History:1"] + if len(s) != 3: + continue + # Please note that title may contain ":". + title = s[2].rsplit(":", maxsplit=1)[0] + titles.add(title) + return list(titles) + + +def get_plotly_graph_objects(system_attrs: dict[str, Any]) -> dict[str, str]: + graph_objects = {} + for title in list_graph_object_ids(system_attrs): + key_prefix = SYSTEM_ATTR_PLOT_DATA + title + ":" + plot_data_attrs = {k: v for k, v in system_attrs.items() if k.startswith(key_prefix)} + graph_objects[title] = concat_plot_data(plot_data_attrs, key_prefix) + return graph_objects + + +def split_plot_data(plot_data_str: str, key_prefix: str) -> dict[str, str]: + plot_data_len = len(plot_data_str) + attrs = {} + for i in range(math.ceil(plot_data_len / SYSTEM_ATTR_MAX_LENGTH)): + start = i * SYSTEM_ATTR_MAX_LENGTH + end = min((i + 1) * SYSTEM_ATTR_MAX_LENGTH, plot_data_len) + attrs[f"{key_prefix}{i}"] = plot_data_str[start:end] + return attrs + + +def concat_plot_data(plot_data_attrs: dict[str, str], key_prefix: str) -> str: + return "".join(plot_data_attrs[f"{key_prefix}{i}"] for i in range(len(plot_data_attrs))) diff --git a/optuna_dashboard/_serializer.py b/optuna_dashboard/_serializer.py index 06b53c42..19acbbd5 100644 --- a/optuna_dashboard/_serializer.py +++ b/optuna_dashboard/_serializer.py @@ -132,6 +132,7 @@ def serialize_study_detail( union: list[tuple[str, BaseDistribution]], union_user_attrs: list[tuple[str, bool]], has_intermediate_values: bool, + plotly_graph_objects: dict[str, str], ) -> dict[str, Any]: serialized: dict[str, Any] = { "name": summary.study_name, @@ -162,6 +163,10 @@ def serialize_study_detail( serialized["form_widgets"] = form_widgets if serialized["is_preferential"]: serialized["preference_history"] = serialize_preference_history(system_attrs) + serialized["plotly_graph_objects"] = [ + {"id": id_, "graph_object": graph_object} + for id_, graph_object in plotly_graph_objects.items() + ] return serialized diff --git a/optuna_dashboard/ts/apiClient.ts b/optuna_dashboard/ts/apiClient.ts index e23fc2ff..e62e0e42 100644 --- a/optuna_dashboard/ts/apiClient.ts +++ b/optuna_dashboard/ts/apiClient.ts @@ -93,6 +93,7 @@ interface StudyDetailResponse { objective_names?: string[] form_widgets?: FormWidgets preference_history?: PreferenceHistoryResponce[] + plotly_graph_objects: PlotlyGraphObject[] } export const getStudyDetailAPI = ( @@ -131,6 +132,7 @@ export const getStudyDetailAPI = ( preference_history: res.data.preference_history?.map( convertPreferenceHistory ), + plotly_graph_objects: res.data.plotly_graph_objects, } }) } diff --git a/optuna_dashboard/ts/components/StudyHistory.tsx b/optuna_dashboard/ts/components/StudyHistory.tsx index b47c557a..5cca3671 100644 --- a/optuna_dashboard/ts/components/StudyHistory.tsx +++ b/optuna_dashboard/ts/components/StudyHistory.tsx @@ -15,6 +15,7 @@ import { GraphIntermediateValues } from "./GraphIntermediateValues" import Grid2 from "@mui/material/Unstable_Grid2" import { DataGrid, DataGridColumn } from "./DataGrid" import { GraphHyperparameterImportance } from "./GraphHyperparameterImportances" +import { UserDefinedPlot } from "./UserDefinedPlot" import { BestTrialsCard } from "./BestTrialsCard" import { useStudyDetailValue, @@ -102,6 +103,17 @@ export const StudyHistory: FC<{ studyId: number }> = ({ studyId }) => { /> + {studyDetail !== null && + studyDetail.plotly_graph_objects.map((go) => ( + + + + ))} {studyDetail !== null && studyDetail.directions.length == 1 && diff --git a/optuna_dashboard/ts/components/UserDefinedPlot.tsx b/optuna_dashboard/ts/components/UserDefinedPlot.tsx new file mode 100644 index 00000000..2a6f98db --- /dev/null +++ b/optuna_dashboard/ts/components/UserDefinedPlot.tsx @@ -0,0 +1,16 @@ +import * as plotly from "plotly.js-dist-min" +import React, { FC, useEffect } from "react" +import { Box } from "@mui/material" + +export const UserDefinedPlot: FC<{ + graphObject: PlotlyGraphObject +}> = ({ graphObject }) => { + const plotDomId = `user-defined-plot:${graphObject.id}` + + useEffect(() => { + const parsed = JSON.parse(graphObject.graph_object) + plotly.react(plotDomId, parsed.data, parsed.layout) + }, [graphObject]) + + return +} diff --git a/optuna_dashboard/ts/types/index.d.ts b/optuna_dashboard/ts/types/index.d.ts index 646d64cf..b7b35797 100644 --- a/optuna_dashboard/ts/types/index.d.ts +++ b/optuna_dashboard/ts/types/index.d.ts @@ -182,6 +182,11 @@ type FormWidgets = widgets: UserAttrFormWidget[] } +type PlotlyGraphObject = { + id: string + graph_object: string +} + type StudyDetail = { id: number name: string @@ -199,6 +204,7 @@ type StudyDetail = { objective_names?: string[] form_widgets?: FormWidgets preference_history?: PreferenceHistory[] + plotly_graph_objects: PlotlyGraphObject[] } type StudyDetails = { diff --git a/pyproject.toml b/pyproject.toml index 0555f701..9b7ce731 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ docs = [ test = [ "coverage", + "plotly", "pytest", "moto[s3]", ] diff --git a/python_tests/test_custom_plot_data.py b/python_tests/test_custom_plot_data.py new file mode 100644 index 00000000..fdf738d7 --- /dev/null +++ b/python_tests/test_custom_plot_data.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from unittest.mock import patch + +import optuna +from optuna_dashboard import _custom_plot_data as custom_plot_data +from optuna_dashboard import save_plotly_graph_object + + +def get_dummy_study() -> optuna.Study: + def objective(trial: optuna.Trial) -> float: + x = trial.suggest_float("x", -100, 100) + y = trial.suggest_categorical("y", [-1, 0, 1]) + return x**2 + y + + study = optuna.create_study() + optuna.logging.set_verbosity(optuna.logging.ERROR) + study.optimize(objective, n_trials=100) + return study + + +def test_save_plotly_graph_object() -> None: + # Save history plot + dummy_study = get_dummy_study() + plot_data = optuna.visualization.plot_optimization_history(dummy_study) + graph_object_id = save_plotly_graph_object(dummy_study, plot_data) + + study_system_attrs = dummy_study._storage.get_study_system_attrs(dummy_study._study_id) + plot_data_dict = custom_plot_data.get_plotly_graph_objects(study_system_attrs) + assert len(plot_data_dict) == 1 + assert plot_data_dict[graph_object_id] == plot_data.to_json() + + # Save parallel coordinate plot + plot_data = optuna.visualization.plot_parallel_coordinate(dummy_study) + graph_object_id = save_plotly_graph_object(dummy_study, plot_data) + + study_system_attrs = dummy_study._storage.get_study_system_attrs(dummy_study._study_id) + plot_data_dict = custom_plot_data.get_plotly_graph_objects(study_system_attrs) + assert len(plot_data_dict) == 2 + assert plot_data_dict[graph_object_id] == plot_data.to_json() + + +def test_update_plotly_graph_object() -> None: + # Save history plot + dummy_study = get_dummy_study() + plot_data = optuna.visualization.plot_optimization_history(dummy_study) + graph_object_id = save_plotly_graph_object(dummy_study, plot_data) + + study_system_attrs = dummy_study._storage.get_study_system_attrs(dummy_study._study_id) + plot_data_dict = custom_plot_data.get_plotly_graph_objects(study_system_attrs) + assert len(plot_data_dict) == 1 + assert plot_data_dict[graph_object_id] == plot_data.to_json() + + # Save parallel coordinate plot + plot_data = optuna.visualization.plot_parallel_coordinate(dummy_study) + graph_object_id = save_plotly_graph_object( + dummy_study, plot_data, graph_object_id=graph_object_id + ) + + study_system_attrs = dummy_study._storage.get_study_system_attrs(dummy_study._study_id) + plot_data_dict = custom_plot_data.get_plotly_graph_objects(study_system_attrs) + assert len(plot_data_dict) == 1 + assert plot_data_dict[graph_object_id] == plot_data.to_json() diff --git a/python_tests/test_serializers.py b/python_tests/test_serializers.py index a90e0de7..72db7b26 100644 --- a/python_tests/test_serializers.py +++ b/python_tests/test_serializers.py @@ -29,7 +29,7 @@ def test_get_study_detail_is_preferential() -> None: assert len(study_summaries) == 1 study_summary = study_summaries[0] - study_detail = serialize_study_detail(study_summary, [], study.trials, [], [], [], False) + study_detail = serialize_study_detail(study_summary, [], study.trials, [], [], [], False, {}) assert study_detail["is_preferential"] @@ -40,7 +40,7 @@ def test_get_study_detail_is_not_preferential() -> None: assert len(study_summaries) == 1 study_summary = study_summaries[0] - study_detail = serialize_study_detail(study_summary, [], study.trials, [], [], [], False) + study_detail = serialize_study_detail(study_summary, [], study.trials, [], [], [], False, {}) assert not study_detail["is_preferential"]