Merge pull request #369 from nabenabe0928/add-plot-target-over-time

Add plot target over time

Author: sawa3030

package/visualization/plot_target_over_time/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Shuhei Watanabe
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/visualization/plot_target_over_time/README.md

@@ -0,0 +1,79 @@
+---
+author: Shuhei Watanabe
+title: Plot Target Over Time
+description: With this visualization module, we can plot the best target value over wall-clock time, averaged across multiple studies with standard error bands.
+tags: [visualization, benchmarking, runtime]
+optuna_versions: [4.8.0]
+license: MIT License
+---
+
+## Abstract
+
+This visualization module enables users to plot the target value over time with standard error bands.
+This module is especially convenient when we use parallel optimization such as [Asynchronous optimization simulation](https://hub.optuna.org/benchmarks/async_opt_simulator).
+
+![Example Using Async Opt](images/async-bench-example.png)
+
+## Class or Function Names
+
+- plot_target_over_time
+
+## Installation
+
+This module requires the following dependencies:
+
+- matplotlib
+- numpy
+
+## APIs
+
+# `plot_target_over_time(study_list, *, ax=None, states=None, target=None, target_direction=None, cumtime_func=None, log_time_scale=True, n_steps=100, color=None, **plot_kwargs)`
+
+- `study_list`: A list of `optuna.Study` objects. Each study is treated as one run, and results are averaged across them.
+- `ax`: A `matplotlib.axes.Axes` object. If not provided, a new figure and axes will be created.
+- `states`: A list of `optuna.trial.TrialState` to include. Defaults to `[TrialState.COMPLETE, TrialState.PRUNED]`.
+- `target`: A callable that takes a `FrozenTrial` and returns a float value. If not provided, `trial.value` is used.
+- `target_direction`: The direction to optimize the target. Required when `target` is specified. Must be `"minimize"`, `"maximize"`, or the corresponding `StudyDirection` enum.
+- `cumtime_func`: A callable that takes a `FrozenTrial` and returns the cumulative time as a float. If not provided, the elapsed time from the first trial start is used.
+- `log_time_scale`: Whether to use a logarithmic time scale for interpolation. Defaults to `True`.
+- `n_steps`: The number of time steps for interpolation. Defaults to `100`.
+- `color`: The color for the plot line and shaded region.
+- `**plot_kwargs`: Additional keyword arguments passed to `ax.plot` (e.g., `label`, `linestyle`).
+
+## Example
+
+```python
+from __future__ import annotations
+
+import optuna
+import optunahub
+
+import matplotlib.pyplot as plt
+
+
+def objective(trial: optuna.Trial) -> float:
+    x = trial.suggest_float("x", -5, 5)
+    y = trial.suggest_float("y", -5, 5)
+    return x**2 + y**2
+
+
+plot_target_over_time = optunahub.load_module("visualization/plot_target_over_time").plot_target_over_time
+_, ax = plt.subplots()
+colors = ["darkred", "black"]
+for sampler, color in zip([optuna.samplers.TPESampler(), optuna.samplers.RandomSampler()], colors):
+    study_list = []
+    for _ in range(5):
+        study = optuna.create_study(sampler=sampler)
+        study.optimize(objective, n_trials=20)
+        study_list.append(study)
+    plot_target_over_time(
+        study_list,
+        ax=ax,
+        color=color,
+        label=sampler.__class__.__name__,
+    )
+
+ax.legend()
+plt.show()
+
+```

package/visualization/plot_target_over_time/__init__.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
+import numpy as np
+import optuna
+from optuna.study import StudyDirection
+from optuna.trial import TrialState
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any
+
+    from matplotlib.axes import Axes
+    from matplotlib.lines import Line2D
+
+
+def _get_values_on_fixed_time_steps(
+    cumtime_list: list[np.ndarray],
+    target_list: list[np.ndarray],
+    log_time_scale: bool,
+    n_steps: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    t_min = np.min(np.stack(cumtime_list))
+    t_max = np.max(np.stack(cumtime_list))
+    if log_time_scale:
+        ts = np.exp(np.linspace(np.log(t_min), np.log(t_max), num=n_steps))
+    else:
+        ts = np.linspace(t_min, t_max, num=n_steps)
+    v_on_grid = []
+    for ct, v in zip(cumtime_list, target_list):
+        i_upper = np.minimum(np.searchsorted(ct, ts, side="left"), v.size - 1)
+        v_on_grid.append(v[i_upper])
+    return ts, np.array(v_on_grid)
+
+
+def _validate(
+    valid_states: tuple[TrialState, ...],
+    states: tuple[TrialState, ...],
+    target: Callable[[optuna.trial.FrozenTrial], float] | None,
+    target_direction: StudyDirection | str | None,
+) -> None:
+    if any(s not in valid_states for s in states):
+        raise ValueError(f"{states=} must be in {valid_states}.")
+    if target_direction is None:
+        if target is not None:
+            raise ValueError("target was specified, but got target_direction=None.")
+    else:
+        if target is None:
+            raise ValueError("target_direction was provided, but got target=None.")
+        if target_direction not in [
+            "minimize",
+            "maximize",
+            StudyDirection.MAXIMIZE,
+            StudyDirection.MINIMIZE,
+        ]:
+            raise ValueError(
+                f"target_direction must be either `minimize` or `maximize` but got {target_direction=}"
+            )
+
+
+def plot_target_over_time(
+    study_list: list[optuna.Study],
+    *,
+    color: str,
+    ax: Axes | None = None,
+    states: tuple[TrialState, ...] | None = None,
+    target: Callable[[optuna.trial.FrozenTrial], float] | None = None,
+    target_direction: optuna.study.StudyDirection | str | None = None,
+    cumtime_func: Callable[[optuna.trial.FrozenTrial], float] | None = None,
+    log_time_scale: bool = True,
+    n_steps: int = 100,
+    **plot_kwargs: Any,
+) -> Line2D:
+    if ax is None:
+        _, ax = plt.subplots()
+
+    valid_states = (TrialState.COMPLETE, TrialState.PRUNED)
+    states = states or valid_states
+    assert states is not None, "Mypy Redefinition."
+    _validate(valid_states, states, target, target_direction)
+
+    target_list = []
+    cumtime_list = []
+    direction = target_direction or study_list[0].direction
+    for study in study_list:
+        trials = study.get_trials(deepcopy=False, states=states)
+        target_vals = np.array([target(t) if target is not None else t.value for t in trials])
+        if cumtime_func is not None:
+            cumtime_list.append(np.array([cumtime_func(t) for t in trials]))
+        else:
+            datetime_start = min(t.datetime_start for t in trials if t.datetime_start is not None)
+            cumtimes = np.array(
+                [
+                    (t.datetime_complete - datetime_start).total_seconds()
+                    for t in trials
+                    if t.datetime_complete is not None
+                ]
+            )
+            cumtime_list.append(cumtimes)
+        order = np.argsort(cumtime_list[-1])
+        cumtime_list[-1] = cumtime_list[-1][order]
+        if direction in ["minimize", StudyDirection.MINIMIZE]:
+            target_list.append(np.minimum.accumulate(target_vals[order]))
+        else:
+            target_list.append(np.maximum.accumulate(target_vals[order]))
+
+    ts, vs = _get_values_on_fixed_time_steps(
+        cumtime_list,
+        target_list,
+        log_time_scale,
+        n_steps,
+    )
+    m = np.mean(vs, axis=0)
+    s = np.std(vs, axis=0) / np.sqrt(len(study_list))
+    (line,) = ax.plot(ts, m, color=color, **plot_kwargs)
+    ax.fill_between(ts, m - s, m + s, color=color, alpha=0.2)
+    return line
+
+
+__all__ = ["plot_target_over_time"]

package/visualization/plot_target_over_time/example.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import optuna
+import optunahub
+
+
+def objective(trial: optuna.Trial) -> float:
+    x = trial.suggest_float("x", -5, 5)
+    y = trial.suggest_float("y", -5, 5)
+    return x**2 + y**2
+
+
+plot_target_over_time = optunahub.load_module(
+    "visualization/plot_target_over_time"
+).plot_target_over_time
+_, ax = plt.subplots()
+colors = ["darkred", "black"]
+for sampler, color in zip([optuna.samplers.TPESampler(), optuna.samplers.RandomSampler()], colors):
+    study_list = []
+    for __ in range(5):
+        study = optuna.create_study(sampler=sampler)
+        study.optimize(objective, n_trials=20)
+        study_list.append(study)
+    plot_target_over_time(
+        study_list,
+        ax=ax,
+        color=color,
+        label=sampler.__class__.__name__,
+    )
+
+ax.legend()
+plt.show()