How to use the optuna.distributions.CategoricalDistribution function in optuna

def _initialize_x0(search_space):
        # type: (Dict[str, BaseDistribution]) -> Dict[str, Any]

        x0 = {}
        for name, distribution in search_space.items():
            if isinstance(distribution, UniformDistribution):
                x0[name] = numpy.mean([distribution.high, distribution.low])
            elif isinstance(distribution, DiscreteUniformDistribution):
                x0[name] = numpy.mean([distribution.high, distribution.low])
            elif isinstance(distribution, IntUniformDistribution):
                x0[name] = int(numpy.mean([distribution.high, distribution.low]))
            elif isinstance(distribution, LogUniformDistribution):
                log_high = math.log(distribution.high)
                log_low = math.log(distribution.low)
                x0[name] = math.exp(numpy.mean([log_high, log_low]))
            elif isinstance(distribution, CategoricalDistribution):
                index = (len(distribution.choices) - 1) // 2
                x0[name] = distribution.choices[index]
            else:
                raise NotImplementedError('The distribution {} is not implemented.'.format(
                    distribution))
        return x0

trial.suggest_loguniform('b', 0.1, 10)
        trial.suggest_discrete_uniform('c', 0, 10, 1)
        trial.suggest_int('d', 0, 10)
        trial.suggest_categorical('e', ['foo', 'bar', 'baz'])

        return 1.0

    study = create_study(storage_init_func())
    study.optimize(objective, n_trials=1)

    assert study.best_trial.distributions == {
        'a': distributions.UniformDistribution(low=0, high=10),
        'b': distributions.LogUniformDistribution(low=0.1, high=10),
        'c': distributions.DiscreteUniformDistribution(low=0, high=10, q=1),
        'd': distributions.IntUniformDistribution(low=0, high=10),
        'e': distributions.CategoricalDistribution(choices=('foo', 'bar', 'baz'))
    }

def test_relative_sampling(storage_mode, cache_mode, comm):
        # type: (str, bool, CommunicatorBase) -> None

        relative_search_space = {
            'x': distributions.UniformDistribution(low=-10, high=10),
            'y': distributions.LogUniformDistribution(low=20, high=30),
            'z': distributions.CategoricalDistribution(choices=(-1.0, 1.0)),
        }
        relative_params = {'x': 1.0, 'y': 25.0, 'z': -1.0}
        sampler = DeterministicRelativeSampler(relative_search_space,  # type: ignore
                                               relative_params)

        with MultiNodeStorageSupplier(storage_mode, cache_mode, comm) as storage:
            study = TestChainerMNStudy._create_shared_study(storage, comm, sampler=sampler)
            mn_study = ChainerMNStudy(study, comm)

            # Invoke optimize.
            n_trials = 20
            func = Func()
            mn_study.optimize(func, n_trials=n_trials)

            # Assert trial counts.
            assert len(mn_study.trials) == n_trials

import pytest

from optuna import distributions
from optuna import type_checking

if type_checking.TYPE_CHECKING:
    from typing import Any  # NOQA
    from typing import Dict  # NOQA
    from typing import List  # NOQA

EXAMPLE_DISTRIBUTIONS = {
    'u': distributions.UniformDistribution(low=1., high=2.),
    'l': distributions.LogUniformDistribution(low=0.001, high=100),
    'du': distributions.DiscreteUniformDistribution(low=1., high=10., q=2.),
    'iu': distributions.IntUniformDistribution(low=1, high=10),
    'c1': distributions.CategoricalDistribution(choices=(2.71, -float('inf'))),
    'c2': distributions.CategoricalDistribution(choices=('Roppongi', 'Azabu'))
}  # type: Dict[str, Any]

EXAMPLE_JSONS = {
    'u': '{"name": "UniformDistribution", "attributes": {"low": 1.0, "high": 2.0}}',
    'l': '{"name": "LogUniformDistribution", "attributes": {"low": 0.001, "high": 100}}',
    'du': '{"name": "DiscreteUniformDistribution",'
    '"attributes": {"low": 1.0, "high": 10.0, "q": 2.0}}',
    'iu': '{"name": "IntUniformDistribution", "attributes": {"low": 1, "high": 10}}',
    'c1': '{"name": "CategoricalDistribution", "attributes": {"choices": [2.71, -Infinity]}}',
    'c2': '{"name": "CategoricalDistribution", "attributes": {"choices": ["Roppongi", "Azabu"]}}'
}


def test_json_to_distribution():
    # type: () -> None

def test_distributions(storage_mode, cache_mode, comm):
        # type: (str, bool, CommunicatorBase) -> None

        with MultiNodeStorageSupplier(storage_mode, cache_mode, comm) as storage:
            study = TestChainerMNStudy._create_shared_study(storage, comm)
            mn_trial = _create_new_chainermn_trial(study, comm)

            mn_trial.suggest_categorical('x', [1])
            assert mn_trial.distributions == {
                'x': distributions.CategoricalDistribution(choices=(1, ))
            }

return self.choices.index(param_value_in_external_repr)

    def single(self):
        # type: () -> bool

        return len(self.choices) == 1

    def _contains(self, param_value_in_internal_repr):
        # type: (float) -> bool

        index = int(param_value_in_internal_repr)
        return 0 <= index and index < len(self.choices)


DISTRIBUTION_CLASSES = (UniformDistribution, LogUniformDistribution, DiscreteUniformDistribution,
                        IntUniformDistribution, CategoricalDistribution)


def json_to_distribution(json_str):
    # type: (str) -> BaseDistribution
    """Deserialize a distribution in JSON format.

    Args:
        json_str: A JSON-serialized distribution.

    Returns:
        A deserialized distribution.
    """

    json_dict = json.loads(json_str)

    if json_dict['name'] == CategoricalDistribution.__name__:

def json_to_distribution(json_str):
    # type: (str) -> BaseDistribution
    """Deserialize a distribution in JSON format.

    Args:
        json_str: A JSON-serialized distribution.

    Returns:
        A deserialized distribution.
    """

    json_dict = json.loads(json_str)

    if json_dict['name'] == CategoricalDistribution.__name__:
        json_dict['attributes']['choices'] = tuple(json_dict['attributes']['choices'])

    for cls in DISTRIBUTION_CLASSES:
        if json_dict['name'] == cls.__name__:
            return cls(**json_dict['attributes'])

    raise ValueError('Unknown distribution class: {}'.format(json_dict['name']))

Note that this method is not supposed to be called by library users.

    Args:
        dist_old: A distribution previously recorded in storage.
        dist_new: A distribution newly added to storage.

    Returns:
        True denotes given distributions are compatible. Otherwise, they are not.
    """

    if dist_old.__class__ != dist_new.__class__:
        raise ValueError('Cannot set different distribution kind to the same parameter name.')

    if not isinstance(dist_old, CategoricalDistribution):
        return
    if not isinstance(dist_new, CategoricalDistribution):
        return
    if dist_old.choices != dist_new.choices:
        raise ValueError(CategoricalDistribution.__name__ +
                         ' does not support dynamic value space.')

>>>     kernel = trial.suggest_categorical('kernel', ['linear', 'poly', 'rbf'])
                >>>     clf = sklearn.svm.SVC(kernel=kernel)
                >>>     ...

        Args:
            name:
                A parameter name.
            choices:
                Candidates of parameter values.

        Returns:
            A suggested value.
        """

        choices = tuple(choices)
        return self._suggest(name, distributions.CategoricalDistribution(choices=choices))

elif isinstance(param_distribution, distributions.DiscreteUniformDistribution):
            return self._sample_discrete_uniform(param_distribution, below_param_values,
                                                 above_param_values)
        elif isinstance(param_distribution, distributions.IntUniformDistribution):
            return self._sample_int(param_distribution, below_param_values, above_param_values)
        elif isinstance(param_distribution, distributions.CategoricalDistribution):
            index = self._sample_categorical_index(param_distribution, below_param_values,
                                                   above_param_values)
            return param_distribution.choices[index]
        else:
            distribution_list = [
                distributions.UniformDistribution.__name__,
                distributions.LogUniformDistribution.__name__,
                distributions.DiscreteUniformDistribution.__name__,
                distributions.IntUniformDistribution.__name__,
                distributions.CategoricalDistribution.__name__
            ]
            raise NotImplementedError("The distribution {} is not implemented. "
                                      "The parameter distribution should be one of the {}".format(
                                          param_distribution, distribution_list))