ax/modelbridge/transforms/trial_as_task.py

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import TYPE_CHECKING, Dict, List, Optional

import numpy as np
from ax.core.observation import ObservationData, ObservationFeatures
from ax.core.parameter import ChoiceParameter, ParameterType
from ax.core.search_space import SearchSpace
from ax.modelbridge.transforms.base import Transform
from ax.models.types import TConfig

if TYPE_CHECKING:
    # import as module to make sphinx-autodoc-typehints happy
    from ax import modelbridge as modelbridge_module  # noqa F401  # pragma: no cover


TRIAL_PARAM = "TRIAL_PARAM"


class TrialAsTask(Transform):
    """Convert trial to one or more task parameters.

    How trial is mapped to parameter is specified with a map like
    {parameter_name: {trial_index: level name}}.
    For example,
    {"trial_param1": {0: "level1", 1: "level1", 2: "level2"},}
    will create choice parameters "trial_param1" with is_task=True.
    Observations with trial 0 or 1 will have "trial_param1" set to "level1",
    and those with trial 2 will have "trial_param1" set to "level2". Multiple
    parameter names and mappings can be specified in this dict.

    The trial level mapping can be specified in config["trial_level_map"]. If
    not specified, defaults to a parameter with a level for every trial index.

    For the reverse transform, if there are multiple mappings in the transform
    the trial will not be set.

    Will raise if trial not specified for every point in the training data.

    Transform is done in-place.
    """

    def __init__(
        self,
        search_space: SearchSpace,
        observation_features: List[ObservationFeatures],
        observation_data: List[ObservationData],
        modelbridge: Optional["modelbridge_module.base.ModelBridge"] = None,
        config: Optional[TConfig] = None,
    ) -> None:
        # Identify values of trial.
        trials = {obsf.trial_index for obsf in observation_features}
        if None in trials:
            raise ValueError(
                "Unable to use trial as task since not all observations have "
                "trial specified."
            )
        # Get trial level map
        if config is not None and "trial_level_map" in config:
            # pyre: Attribute `trial_level_map` declared in class `ax.
            # pyre: modelbridge.transforms.trial_as_task.TrialAsTask` has type
            # pyre: `Dict[str, Dict[int, str]]` but is used as type `typing.
            # pyre-fixme[8]: Union[float, int, str]`.
            self.trial_level_map: Dict[str, Dict[int, str]] = config["trial_level_map"]
            # Validate
            for _p_name, level_dict in self.trial_level_map.items():
                # Check that trials match those in data
                level_map = set(level_dict.keys())
                if not trials.issubset(level_map):
                    raise ValueError(
                        f"Not all trials in data ({trials}) contained "
                        f"in trial level map for {_p_name} ({level_map})"
                    )
        else:
            # Set TRIAL_PARAM for each trial to the corresponding trial_index.
            # pyre-fixme[6]: Expected `Union[bytes, str, typing.SupportsInt]` for
            #  1st param but got `Optional[np.int64]`.
            self.trial_level_map = {TRIAL_PARAM: {int(b): str(b) for b in trials}}
        if len(self.trial_level_map) == 1:
            level_dict = next(iter(self.trial_level_map.values()))
            self.inverse_map: Optional[Dict[str, int]] = {
                v: k for k, v in level_dict.items()
            }
        else:
            self.inverse_map = None

    def transform_observation_features(
        self, observation_features: List[ObservationFeatures]
    ) -> List[ObservationFeatures]:
        for obsf in observation_features:
            if obsf.trial_index is not None:
                for p_name, level_dict in self.trial_level_map.items():
                    # pyre-fixme[6]: Expected `Union[bytes, str,
                    #  typing.SupportsInt]` for 1st param but got `Optional[np.int64]`.
                    obsf.parameters[p_name] = level_dict[int(obsf.trial_index)]
                obsf.trial_index = None
        return observation_features

    def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
        for p_name, level_dict in self.trial_level_map.items():
            level_values = sorted(set(level_dict.values()))
            if len(level_values) < 2:
                details = (
                    f"only 1 found: {level_values}" if level_values else "none found"
                )
                raise ValueError(
                    f"TrialAsTask transform expects 2+ task params, {details}"
                )
            trial_param = ChoiceParameter(
                name=p_name,
                parameter_type=ParameterType.STRING,
                # Expected `List[Optional[typing.Union[bool, float, str]]]` for 4th
                # parameter `values` to call
                # `ax.core.parameter.ChoiceParameter.__init__` but got
                # `List[str]`.
                # pyre-fixme[6]:
                values=level_values,
                is_ordered=False,
                is_task=True,
                sort_values=True,
            )
            search_space.add_parameter(trial_param)
        return search_space

    def untransform_observation_features(
        self, observation_features: List[ObservationFeatures]
    ) -> List[ObservationFeatures]:
        for obsf in observation_features:
            for p_name in self.trial_level_map:
                pval = obsf.parameters.pop(p_name)
            if self.inverse_map is not None:
                # pyre-fixme[61]: `pval` may not be initialized here.
                obsf.trial_index = np.int64(self.inverse_map[pval])
        return observation_features