optuna

1
from __future__ import annotations
2

3
import pandas as pd
4
import pytest
5

6
from optuna import create_study
7
from optuna import create_trial
8
from optuna import Trial
9
from optuna.testing.storages import STORAGE_MODES
10
from optuna.testing.storages import StorageSupplier
11
from optuna.trial import TrialState
12

13

14
def test_study_trials_dataframe_with_no_trials() -> None:
15
    study_with_no_trials = create_study()
16
    trials_df = study_with_no_trials.trials_dataframe()
17
    assert trials_df.empty
18

19

20
@pytest.mark.parametrize("storage_mode", STORAGE_MODES)
21
@pytest.mark.parametrize(
22
    "attrs",
23
    [
24
        (
25
            "number",
26
            "value",
27
            "datetime_start",
28
            "datetime_complete",
29
            "params",
30
            "user_attrs",
31
            "system_attrs",
32
            "state",
33
        ),
34
        (
35
            "number",
36
            "value",
37
            "datetime_start",
38
            "datetime_complete",
39
            "duration",
40
            "params",
41
            "user_attrs",
42
            "system_attrs",
43
            "state",
44
            "intermediate_values",
45
            "_trial_id",
46
            "distributions",
47
        ),
48
    ],
49
)
50
@pytest.mark.parametrize("multi_index", [True, False])
51
def test_trials_dataframe(storage_mode: str, attrs: tuple[str, ...], multi_index: bool) -> None:
52
    def f(trial: Trial) -> float:
53
        x = trial.suggest_int("x", 1, 1)
54
        y = trial.suggest_categorical("y", (2.5,))
55
        trial.set_user_attr("train_loss", 3)
56
        trial.storage.set_trial_system_attr(trial._trial_id, "foo", "bar")
57
        value = x + y  # 3.5
58

59
        # Test reported intermediate values, although it in practice is not "intermediate".
60
        trial.report(value, step=0)
61

62
        return value
63

64
    with StorageSupplier(storage_mode) as storage:
65
        study = create_study(storage=storage)
66
        study.optimize(f, n_trials=3)
67
        df = study.trials_dataframe(attrs=attrs, multi_index=multi_index)
68
        # Change index to access rows via trial number.
69
        if multi_index:
70
            df.set_index(("number", ""), inplace=True, drop=False)
71
        else:
72
            df.set_index("number", inplace=True, drop=False)
73
        assert len(df) == 3
74

75
        # Number columns are as follows (total of 13):
76
        #   non-nested: 6 (number, value, state, datetime_start, datetime_complete, duration)
77
        #   params: 2
78
        #   distributions: 2
79
        #   user_attrs: 1
80
        #   system_attrs: 1
81
        #   intermediate_values: 1
82
        expected_n_columns = len(attrs)
83
        if "params" in attrs:
84
            expected_n_columns += 1
85
        if "distributions" in attrs:
86
            expected_n_columns += 1
87
        assert len(df.columns) == expected_n_columns
88

89
        for i in range(3):
90
            assert df.number[i] == i
91
            assert df.state[i] == "COMPLETE"
92
            assert df.value[i] == 3.5
93
            assert isinstance(df.datetime_start[i], pd.Timestamp)
94
            assert isinstance(df.datetime_complete[i], pd.Timestamp)
95

96
            if multi_index:
97
                if "distributions" in attrs:
98
                    assert ("distributions", "x") in df.columns
99
                    assert ("distributions", "y") in df.columns
100
                if "_trial_id" in attrs:
101
                    assert ("trial_id", "") in df.columns  # trial_id depends on other tests.
102
                if "duration" in attrs:
103
                    assert ("duration", "") in df.columns
104

105
                assert df.params.x[i] == 1
106
                assert df.params.y[i] == 2.5
107
                assert df.user_attrs.train_loss[i] == 3
108
                assert df.system_attrs.foo[i] == "bar"
109
            else:
110
                if "distributions" in attrs:
111
                    assert "distributions_x" in df.columns
112
                    assert "distributions_y" in df.columns
113
                if "_trial_id" in attrs:
114
                    assert "trial_id" in df.columns  # trial_id depends on other tests.
115
                if "duration" in attrs:
116
                    assert "duration" in df.columns
117

118
                assert df.params_x[i] == 1
119
                assert df.params_y[i] == 2.5
120
                assert df.user_attrs_train_loss[i] == 3
121
                assert df.system_attrs_foo[i] == "bar"
122

123

124
@pytest.mark.parametrize("storage_mode", STORAGE_MODES)
125
def test_trials_dataframe_with_failure(storage_mode: str) -> None:
126
    def f(trial: Trial) -> float:
127
        x = trial.suggest_int("x", 1, 1)
128
        y = trial.suggest_categorical("y", (2.5,))
129
        trial.set_user_attr("train_loss", 3)
130
        raise ValueError()
131
        return x + y  # 3.5
132

133
    with StorageSupplier(storage_mode) as storage:
134
        study = create_study(storage=storage)
135
        study.optimize(f, n_trials=3, catch=(ValueError,))
136
        df = study.trials_dataframe()
137
        # Change index to access rows via trial number.
138
        df.set_index("number", inplace=True, drop=False)
139
        assert len(df) == 3
140
        # non-nested: 6, params: 2, user_attrs: 1 system_attrs: 0
141
        assert len(df.columns) == 9
142
        for i in range(3):
143
            assert df.number[i] == i
144
            assert df.state[i] == "FAIL"
145
            assert df.value[i] is None
146
            assert isinstance(df.datetime_start[i], pd.Timestamp)
147
            assert isinstance(df.datetime_complete[i], pd.Timestamp)
148
            assert isinstance(df.duration[i], pd.Timedelta)
149
            assert df.params_x[i] == 1
150
            assert df.params_y[i] == 2.5
151
            assert df.user_attrs_train_loss[i] == 3
152

153

154
@pytest.mark.parametrize("attrs", [("value",), ("values",)])
155
@pytest.mark.parametrize("multi_index", [True, False])
156
def test_trials_dataframe_with_multi_objective_optimization(
157
    attrs: tuple[str, ...], multi_index: bool
158
) -> None:
159
    def f(trial: Trial) -> tuple[float, float]:
160
        x = trial.suggest_float("x", 1, 1)
161
        y = trial.suggest_float("y", 2, 2)
162

163
        return x + y, x**2 + y**2  # 3, 5
164

165
    # without set_metric_names()
166
    study = create_study(directions=["minimize", "maximize"])
167
    study.optimize(f, n_trials=1)
168
    df = study.trials_dataframe(attrs=attrs, multi_index=multi_index)
169
    if multi_index:
170
        assert df.get("values")[0][0] == 3
171
        assert df.get("values")[1][0] == 5
172
    else:
173
        assert df.values_0[0] == 3
174
        assert df.values_1[0] == 5
175

176
    # with set_metric_names()
177
    study.set_metric_names(["v0", "v1"])
178
    df = study.trials_dataframe(attrs=attrs, multi_index=multi_index)
179
    if multi_index:
180
        assert df.get("values")["v0"][0] == 3
181
        assert df.get("values")["v1"][0] == 5
182
    else:
183
        assert df.get("values_v0")[0] == 3
184
        assert df.get("values_v1")[0] == 5
185

186

187
@pytest.mark.parametrize("attrs", [("value",), ("values",)])
188
@pytest.mark.parametrize("multi_index", [True, False])
189
def test_trials_dataframe_with_multi_objective_optimization_with_fail_and_pruned(
190
    attrs: tuple[str, ...], multi_index: bool
191
) -> None:
192
    study = create_study(directions=["minimize", "maximize"])
193
    study.add_trial(create_trial(state=TrialState.FAIL))
194
    study.add_trial(create_trial(state=TrialState.PRUNED))
195
    df = study.trials_dataframe(attrs=attrs, multi_index=multi_index)
196

197
    # without set_metric_names()
198
    if multi_index:
199
        for i in range(2):
200
            assert df.get("values")[0][i] is None
201
            assert df.get("values")[1][i] is None
202
    else:
203
        for i in range(2):
204
            assert df.values_0[i] is None
205
            assert df.values_1[i] is None
206

207
    # with set_metric_names()
208
    study.set_metric_names(["v0", "v1"])
209
    df = study.trials_dataframe(attrs=attrs, multi_index=multi_index)
210
    if multi_index:
211
        assert df.get("values")["v0"][0] is None
212
        assert df.get("values")["v1"][0] is None
213
    else:
214
        assert df.get("values_v0")[0] is None
215
        assert df.get("values_v1")[0] is None
216