Implement OneHot.transform_experiment_data (facebook#3886)

Summary:
Pull Request resolved: facebook#3886

As titled. Supports transforming `ExperimentData` with `OneHot` transform.

Background: As part of the larger refactor, we will be using `ExperimentData` in place of `list[Observation]` within the `Adapter`.
- The transforms will be initialized using `ExperimentData`. The `observations` input to the constructors may be deprecated once the use cases are updated.
- The training data for `Adapter` will be represented with `ExperimentData` and will be transformed using `transform_experiment_data`.
- For misc input / output to various `Adapter` and other methods, the `Observation / ObservationFeatures / ObservationData` objects will remain. To support these, we will retain the existing transform methods that service these objects.
- Since `ExperimentData` is not planned to be used as an output of user facing methods, we do not need to untransform it. We are not planning to implement`untransform_experiment_data`.

Reviewed By: esantorella

Differential Revision: D76074577

fbshipit-source-id: 07f17b5f91a7e1ce5c6d9eb91dd034bec17d6f89

Author: saitcakmak

ax/adapter/transforms/one_hot.py

@@ -9,6 +9,7 @@
 from typing import Optional, TYPE_CHECKING
 
 import numpy as np
+import pandas as pd
 from ax.adapter.data_utils import ExperimentData
 from ax.adapter.transforms.base import Transform
 from ax.adapter.transforms.rounding import randomized_onehot_round, strict_onehot_round
@@ -207,3 +208,35 @@ def untransform_observation_features(
                 )
                 obsf.parameters[p_name] = val
         return observation_features
+
+    def transform_experiment_data(
+        self, experiment_data: ExperimentData
+    ) -> ExperimentData:
+        arm_data = experiment_data.arm_data
+        for p_name, values in self.encoded_values.items():
+            # First, replace values with 0, 1, 2, so that column names are as expected.
+            arm_data = arm_data.replace(
+                to_replace={p_name: {v: i for i, v in enumerate(values)}}
+            ).astype({p_name: int})
+
+            if len(values) == 2:
+                # Handle the special case. Only need to rename the column.
+                arm_data = arm_data.rename(columns={p_name: p_name + OH_PARAM_INFIX})
+            else:
+                # Use get_dummies to one-hot encode the column.
+                arm_data = pd.get_dummies(
+                    arm_data,
+                    columns=[p_name],
+                    prefix=p_name + OH_PARAM_INFIX,
+                    # Could be int, but using float to match the parameter type.
+                    dtype=float,
+                )
+                # Make sure all expected columns are present, even if there is no
+                # corresponding value in the data.
+                for i in range(len(values)):
+                    if f"{p_name}{OH_PARAM_INFIX}_{i}" not in arm_data:
+                        arm_data[f"{p_name}{OH_PARAM_INFIX}_{i}"] = 0.0
+
+        return ExperimentData(
+            arm_data=arm_data, observation_data=experiment_data.observation_data
+        )

ax/adapter/transforms/tests/test_one_hot_transform.py

@@ -8,14 +8,22 @@
 
 from copy import deepcopy
 
-from ax.adapter.transforms.one_hot import OH_PARAM_INFIX, OneHot
+from ax.adapter.base import DataLoaderConfig
+
+from ax.adapter.data_utils import extract_experiment_data
 
+from ax.adapter.transforms.one_hot import OH_PARAM_INFIX, OneHot
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import RobustSearchSpace, SearchSpace
 from ax.utils.common.testutils import TestCase
-from ax.utils.testing.core_stubs import get_robust_search_space
+from ax.utils.testing.core_stubs import (
+    get_experiment_with_observations,
+    get_robust_search_space,
+)
+from pandas import DataFrame
+from pandas.testing import assert_frame_equal
 
 
 class OneHotTransformTest(TestCase):
@@ -34,9 +42,7 @@ def setUp(self) -> None:
                     "b", parameter_type=ParameterType.STRING, values=["a", "b", "c"]
                 ),
                 ChoiceParameter(
-                    "c",
-                    parameter_type=ParameterType.BOOL,
-                    values=[True, False],
+                    "c", parameter_type=ParameterType.BOOL, values=[True, False]
                 ),
                 ChoiceParameter(
                     "d",
@@ -49,13 +55,9 @@ def setUp(self) -> None:
                 ParameterConstraint(constraint_dict={"x": -0.5, "a": 1}, bound=0.5)
             ],
         )
-        self.t = OneHot(
-            search_space=self.search_space,
-            observations=[],
-        )
+        self.t = OneHot(search_space=self.search_space)
         self.t2 = OneHot(
             search_space=self.search_space,
-            observations=[],
             config={"rounding": "randomized"},
         )
 
@@ -255,3 +257,55 @@ def test_heterogeneous_search_space(self) -> None:
         ]
         untf_obs = self.t.untransform_observation_features(obs_ft)
         self.assertFalse(any(obs.parameters.get("b") == "b" for obs in untf_obs))
+
+    def test_transform_experiment_data(self) -> None:
+        parameterizations = [
+            {"x": 2.2, "a": 2, "b": "b", "c": False, "d": 10.0},
+            {"x": 1.2, "a": 2, "b": "a", "c": False, "d": 100.0},
+        ]
+        experiment = get_experiment_with_observations(
+            observations=[[1.0], [2.0]],
+            search_space=self.search_space,
+            parameterizations=parameterizations,
+        )
+        experiment_data = extract_experiment_data(
+            experiment=experiment, data_loader_config=DataLoaderConfig()
+        )
+        transformed_data = self.t.transform_experiment_data(
+            experiment_data=deepcopy(experiment_data)
+        )
+
+        # Check that only "b" has been transformed and column names are as expected.
+        base_columns = ["x", "a", "c", "d", "metadata"]
+        transformed_columns = [
+            "b" + OH_PARAM_INFIX + "_0",
+            "b" + OH_PARAM_INFIX + "_1",
+            "b" + OH_PARAM_INFIX + "_2",
+        ]
+        self.assertEqual(
+            set(transformed_data.arm_data),
+            {*base_columns, *transformed_columns},
+        )
+
+        # Untransformed columns are same as before.
+        assert_frame_equal(
+            transformed_data.arm_data[base_columns],
+            experiment_data.arm_data[base_columns],
+        )
+        # Observation data is unchanged.
+        assert_frame_equal(
+            transformed_data.observation_data, experiment_data.observation_data
+        )
+
+        # Transformed columns have correct values.
+        expected_columns = DataFrame(
+            index=transformed_data.arm_data.index,
+            data=[
+                [0.0, 1.0, 0.0],
+                [1.0, 0.0, 0.0],
+            ],
+            columns=transformed_columns,
+        )
+        assert_frame_equal(
+            transformed_data.arm_data[transformed_columns], expected_columns
+        )