optuna-dashboard/python_tests/test_cached_extra_study_property.py

from typing import Dict
from typing import List
from unittest import TestCase
import warnings

import optuna
from optuna import create_trial
from optuna.distributions import BaseDistribution
from optuna.distributions import UniformDistribution
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": UniformDistribution(low=0, high=10),
                "x1": UniformDistribution(low=0, high=10),
            },
            {
                "x0": UniformDistribution(low=0, high=10),
                "x1": UniformDistribution(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), 2)
        self.assertEqual(len(cached_extra_study_property.union), 2)

    def test_different_distributions(self) -> None:
        distributions: List[Dict[str, BaseDistribution]] = [
            {
                "x0": UniformDistribution(low=0, high=10),
                "x1": UniformDistribution(low=0, high=10),
            },
            {
                "x0": UniformDistribution(low=0, high=5),
                "x1": UniformDistribution(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), 1)
        self.assertEqual(len(cached_extra_study_property.union), 3)

    def test_dynamic_search_space(self) -> None:
        distributions: List[Dict[str, BaseDistribution]] = [
            {
                "x0": UniformDistribution(low=0, high=10),
                "x1": UniformDistribution(low=0, high=10),
            },
            {
                "x0": UniformDistribution(low=0, high=5),
            },
            {
                "x0": UniformDistribution(low=0, high=10),
                "x1": UniformDistribution(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), 0)
        self.assertEqual(len(cached_extra_study_property.union), 3)

    def test_contains_failed_trials(self) -> None:
        distributions = {
            "x0": UniformDistribution(low=0, high=10),
            "x1": UniformDistribution(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=0, 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), 2)
        self.assertEqual(len(cached_extra_study_property.union), 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": UniformDistribution(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": UniformDistribution(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": UniformDistribution(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 = {
            "x0": UniformDistribution(low=0, high=10),
            "x1": UniformDistribution(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=0,
                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)