optuna-dashboard/python_tests/test_cached_extra_study_property.py

from __future__ import annotations

from typing import Any
from unittest import TestCase
import warnings

import numpy as np
import optuna
from optuna import create_trial
from optuna.distributions import BaseDistribution
from optuna.distributions import FloatDistribution
from optuna.exceptions import ExperimentalWarning
from optuna.trial import TrialState
from optuna_dashboard._cached_extra_study_property import _CachedExtraStudyProperty


class _CachedExtraStudyPropertySearchSpaceTestCase(TestCase):
    def setUp(self) -> None:
        optuna.logging.set_verbosity(optuna.logging.ERROR)
        warnings.simplefilter("ignore", category=ExperimentalWarning)

    def test_same_distributions(self) -> None:
        distributions: list[dict[str, BaseDistribution]] = [
            {
                "x0": FloatDistribution(low=0, high=10),
                "x1": FloatDistribution(low=0, high=10),
            },
            {
                "x0": FloatDistribution(low=0, high=10),
                "x1": FloatDistribution(low=0, high=10),
            },
        ]
        params = [
            {
                "x0": 0.5,
                "x1": 0.5,
            },
            {
                "x0": 0.5,
                "x1": 0.5,
            },
        ]
        trials = [
            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
            for d, p in zip(distributions, params)
        ]
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)

        self.assertEqual(len(cached_extra_study_property.intersection_search_space), 2)
        self.assertEqual(len(cached_extra_study_property.union_search_space), 2)

    def test_different_distributions(self) -> None:
        distributions: list[dict[str, BaseDistribution]] = [
            {
                "x0": FloatDistribution(low=0, high=10),
                "x1": FloatDistribution(low=0, high=10),
            },
            {
                "x0": FloatDistribution(low=0, high=5),
                "x1": FloatDistribution(low=0, high=10),
            },
        ]
        params = [
            {
                "x0": 0.5,
                "x1": 0.5,
            },
            {
                "x0": 0.5,
                "x1": 0.5,
            },
        ]
        trials = [
            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
            for d, p in zip(distributions, params)
        ]
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)

        self.assertEqual(len(cached_extra_study_property.intersection_search_space), 1)
        self.assertEqual(len(cached_extra_study_property.union_search_space), 3)

    def test_dynamic_search_space(self) -> None:
        distributions: list[dict[str, BaseDistribution]] = [
            {
                "x0": FloatDistribution(low=0, high=10),
                "x1": FloatDistribution(low=0, high=10),
            },
            {
                "x0": FloatDistribution(low=0, high=5),
            },
            {
                "x0": FloatDistribution(low=0, high=10),
                "x1": FloatDistribution(low=0, high=10),
            },
        ]
        params = [
            {
                "x0": 0.5,
                "x1": 0.5,
            },
            {
                "x0": 0.5,
            },
            {
                "x0": 0.5,
                "x1": 0.5,
            },
        ]
        trials = [
            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
            for d, p in zip(distributions, params)
        ]

        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)

        self.assertEqual(len(cached_extra_study_property.intersection_search_space), 0)
        self.assertEqual(len(cached_extra_study_property.union_search_space), 3)

    def test_contains_failed_trials(self) -> None:
        distributions: dict[str, BaseDistribution] = {
            "x0": FloatDistribution(low=0, high=10),
            "x1": FloatDistribution(low=0, high=10),
        }
        params = {
            "x0": 0.5,
            "x1": 0.5,
        }
        trials = [
            create_trial(
                state=TrialState.COMPLETE, value=0, distributions=distributions, params=params
            ),
            create_trial(state=TrialState.FAIL, value=None, distributions={}, params={}),
            create_trial(
                state=TrialState.COMPLETE, value=0, distributions=distributions, params=params
            ),
        ]
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)

        self.assertEqual(len(cached_extra_study_property.intersection_search_space), 2)
        self.assertEqual(len(cached_extra_study_property.union_search_space), 2)


class _CachedExtraStudyPropertyIntermediateTestCase(TestCase):
    def setUp(self) -> None:
        optuna.logging.set_verbosity(optuna.logging.ERROR)
        warnings.simplefilter("ignore", category=ExperimentalWarning)

    def test_no_intermediate_value(self) -> None:
        intermediate_values: list[dict] = [
            {},
            {},
        ]
        trials = [
            create_trial(
                state=TrialState.COMPLETE,
                value=0,
                distributions={"x0": FloatDistribution(low=0, high=10)},
                intermediate_values=iv,
                params={"x0": 0.5},
            )
            for iv in intermediate_values
        ]
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)
        self.assertFalse(cached_extra_study_property.has_intermediate_values)

    def test_some_trials_has_no_intermediate_value(self) -> None:
        intermediate_values: list[dict] = [
            {0: 0.3, 1: 1.2},
            {},
            {0: 0.3, 1: 1.2},
        ]
        trials = [
            create_trial(
                state=TrialState.COMPLETE,
                value=0,
                distributions={"x0": FloatDistribution(low=0, high=10)},
                intermediate_values=iv,
                params={"x0": 0.5},
            )
            for iv in intermediate_values
        ]
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)
        self.assertTrue(cached_extra_study_property.has_intermediate_values)

    def test_all_trials_has_intermediate_value(self) -> None:
        intermediate_values: list[dict] = [{0: 0.3, 1: 1.2}, {0: 0.3, 1: 1.2}]
        trials = [
            create_trial(
                state=TrialState.COMPLETE,
                value=0,
                distributions={"x0": FloatDistribution(low=0, high=10)},
                intermediate_values=iv,
                params={"x0": 0.5},
            )
            for iv in intermediate_values
        ]
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)
        self.assertTrue(cached_extra_study_property.has_intermediate_values)

    def test_no_trials(self) -> None:
        trials: list = []
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)
        self.assertFalse(cached_extra_study_property.has_intermediate_values)


class _CachedExtraStudyPropertyUserAttrs(TestCase):
    def setUp(self) -> None:
        optuna.logging.set_verbosity(optuna.logging.ERROR)
        warnings.simplefilter("ignore", category=ExperimentalWarning)

    def test_contains_failed_trials(self) -> None:
        distributions: dict[str, BaseDistribution] = {
            "x0": FloatDistribution(low=0, high=10),
            "x1": FloatDistribution(low=0, high=10),
        }
        params = {
            "x0": 0.5,
            "x1": 0.5,
        }
        trials = [
            create_trial(
                state=TrialState.COMPLETE,
                value=0,
                distributions=distributions,
                params=params,
                user_attrs={"foo": "foo"},
            ),
            create_trial(
                state=TrialState.FAIL,
                value=None,
                distributions={},
                params={},
                user_attrs={"bar": "bar"},
            ),
            create_trial(
                state=TrialState.COMPLETE,
                value=0,
                distributions=distributions,
                params=params,
                user_attrs={"baz": "baz"},
            ),
        ]
        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)

        self.assertEqual(len(cached_extra_study_property.union_user_attrs), 3)

    def test_infer_sortable(self) -> None:
        user_attrs_list: list[dict[str, Any]] = [
            {
                "a": 1,
                "b": 1,
                "c": 1,
                "d": "a",
                "e": 1,
                "f": True,
                "g": np.float128(1.1),
                "h": np.int64(2),
            },
            {"a": 2, "b": "a", "c": "a", "d": "a"},
            {"a": 3, "b": None, "c": 3, "d": "a", "e": 3},
        ]
        expected = {
            "a": True,
            "b": False,
            "c": False,
            "d": False,
            "e": True,
            "f": False,
            "g": True,
            "h": True,
        }

        trials = []
        for user_attrs in user_attrs_list:
            trials.append(
                create_trial(
                    state=TrialState.COMPLETE,
                    value=0,
                    distributions={
                        "x0": FloatDistribution(low=0, high=10),
                        "x1": FloatDistribution(low=0, high=10),
                    },
                    params={"x0": 0.5, "x1": 0.5},
                    user_attrs=user_attrs,
                )
            )

        cached_extra_study_property = _CachedExtraStudyProperty()
        cached_extra_study_property.update(trials)
        actual = {k: v for k, v in cached_extra_study_property.union_user_attrs}
        self.assertEqual(actual, expected)