Implement UnitX.transform_experiment_data (facebook#3888)

Summary:
Pull Request resolved: facebook#3888

As titled. Supports transforming `ExperimentData` with `UnitX` 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: D76085733

fbshipit-source-id: 40a07675cfc4252f9d5699423e6e2bfea5295378

Author: saitcakmak

ax/adapter/transforms/tests/test_unit_x_transform.py

@@ -9,28 +9,27 @@
 from copy import deepcopy
 
 import numpy as np
+from ax.adapter.base import DataLoaderConfig
+from ax.adapter.data_utils import extract_experiment_data
 from ax.adapter.transforms.unit_x import UnitX
 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.exceptions.core import UnsupportedError, UserInputError
 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
 from pyre_extensions import assert_is_instance
 
 
 class UnitXTransformTest(TestCase):
-    transform_class = UnitX
-    # pyre-fixme[4]: Attribute must be annotated.
-    expected_c_dicts = [{"x": -1.0, "y": 1.0}, {"x": -1.0, "a": 1.0}]
-    expected_c_bounds = [0.0, 1.0]
-
     def setUp(self) -> None:
         super().setUp()
-        self.target_lb = self.transform_class.target_lb
-        self.target_range = self.transform_class.target_range
-        self.target_ub = self.target_lb + self.target_range
         self.search_space = SearchSpace(
             parameters=[
                 RangeParameter(
@@ -56,10 +55,7 @@ def setUp(self) -> None:
                 ParameterConstraint(constraint_dict={"x": -0.5, "a": 1}, bound=0.5),
             ],
         )
-        self.t = self.transform_class(
-            search_space=self.search_space,
-            observations=[],
-        )
+        self.t = UnitX(search_space=self.search_space)
         self.search_space_with_target = SearchSpace(
             parameters=[
                 RangeParameter(
@@ -86,13 +82,7 @@ def test_TransformObservationFeatures(self) -> None:
             obs_ft2,
             [
                 ObservationFeatures(
-                    parameters={
-                        "x": self.target_lb + self.target_range / 2.0,
-                        "y": 1.0,
-                        "z": 2,
-                        "a": 2,
-                        "b": "b",
-                    }
+                    parameters={"x": 0.5, "y": 1.0, "z": 2, "a": 2, "b": "b"}
                 )
             ],
         )
@@ -103,7 +93,7 @@ def test_TransformObservationFeatures(self) -> None:
         obs_ft3 = self.t.transform_observation_features(obs_ft3)
         self.assertEqual(
             obs_ft3[0],
-            ObservationFeatures(parameters={"x": self.target_ub, "z": 2}),
+            ObservationFeatures(parameters={"x": 1.0, "z": 2}),
         )
         obs_ft5 = self.t.transform_observation_features([ObservationFeatures({})])
         self.assertEqual(obs_ft5[0], ObservationFeatures({}))
@@ -114,31 +104,35 @@ def test_TransformSearchSpace(self) -> None:
 
         # Parameters transformed
         true_bounds = {
-            "x": (self.target_lb, 1.0),
-            "y": (self.target_lb, 1.0),
+            "x": (0.0, 1.0),
+            "y": (0.0, 1.0),
             "z": (1.0, 2.0),
             "a": (1.0, 2.0),
         }
         for p_name, (l, u) in true_bounds.items():
-            self.assertEqual(ss2.parameters[p_name].lower, l)
-            self.assertEqual(ss2.parameters[p_name].upper, u)
-        self.assertEqual(ss2.parameters["b"].values, ["a", "b", "c"])
+            self.assertEqual(
+                assert_is_instance(ss2.parameters[p_name], RangeParameter).lower, l
+            )
+            self.assertEqual(
+                assert_is_instance(ss2.parameters[p_name], RangeParameter).upper, u
+            )
+        self.assertEqual(
+            assert_is_instance(ss2.parameters["b"], ChoiceParameter).values,
+            ["a", "b", "c"],
+        )
         self.assertEqual(len(ss2.parameters), 5)
         # Constraints transformed
         self.assertEqual(
-            ss2.parameter_constraints[0].constraint_dict, self.expected_c_dicts[0]
+            ss2.parameter_constraints[0].constraint_dict, {"x": -1.0, "y": 1.0}
         )
-        self.assertEqual(ss2.parameter_constraints[0].bound, self.expected_c_bounds[0])
+        self.assertEqual(ss2.parameter_constraints[0].bound, 0.0)
         self.assertEqual(
-            ss2.parameter_constraints[1].constraint_dict, self.expected_c_dicts[1]
+            ss2.parameter_constraints[1].constraint_dict, {"x": -1.0, "a": 1.0}
         )
-        self.assertEqual(ss2.parameter_constraints[1].bound, self.expected_c_bounds[1])
+        self.assertEqual(ss2.parameter_constraints[1].bound, 1.0)
 
         # Test transform of target value
-        t = self.transform_class(
-            search_space=self.search_space_with_target,
-            observations=[],
-        )
+        t = UnitX(search_space=self.search_space_with_target)
         t.transform_search_space(self.search_space_with_target)
         self.assertEqual(
             self.search_space_with_target.parameters["x"].target_value, 1.0
@@ -175,14 +169,8 @@ def test_TransformNewSearchSpace(self) -> None:
         self.t.transform_search_space(new_ss)
         # Parameters transformed
         true_bounds = {
-            "x": [
-                0.25 * self.target_range + self.target_lb,
-                0.5 * self.target_range + self.target_lb,
-            ],
-            "y": [
-                0.25 * self.target_range + self.target_lb,
-                1.0 * self.target_range + self.target_lb,
-            ],
+            "x": [0.25, 0.5],
+            "y": [0.25, 1.0],
             "z": [1.0, 1.5],
             "a": [0, 2],
         }
@@ -197,23 +185,16 @@ def test_TransformNewSearchSpace(self) -> None:
         self.assertEqual(len(new_ss.parameters), 5)
         # # Constraints transformed
         self.assertEqual(
-            new_ss.parameter_constraints[0].constraint_dict, self.expected_c_dicts[0]
-        )
-        self.assertEqual(
-            new_ss.parameter_constraints[0].bound, self.expected_c_bounds[0]
-        )
-        self.assertEqual(
-            new_ss.parameter_constraints[1].constraint_dict, self.expected_c_dicts[1]
+            new_ss.parameter_constraints[0].constraint_dict, {"x": -1.0, "y": 1.0}
         )
+        self.assertEqual(new_ss.parameter_constraints[0].bound, 0.0)
         self.assertEqual(
-            new_ss.parameter_constraints[1].bound, self.expected_c_bounds[1]
+            new_ss.parameter_constraints[1].constraint_dict, {"x": -1.0, "a": 1.0}
         )
+        self.assertEqual(new_ss.parameter_constraints[1].bound, 1.0)
 
         # Test transform of target value
-        t = self.transform_class(
-            search_space=self.search_space_with_target,
-            observations=[],
-        )
+        t = UnitX(search_space=self.search_space_with_target)
         new_search_space_with_target = SearchSpace(
             parameters=[
                 RangeParameter(
@@ -227,50 +208,30 @@ def test_TransformNewSearchSpace(self) -> None:
             ]
         )
         t.transform_search_space(new_search_space_with_target)
-        self.assertEqual(
-            new_search_space_with_target.parameters["x"].target_value,
-            0.5 * self.target_range + self.target_lb,
-        )
+        self.assertEqual(new_search_space_with_target.parameters["x"].target_value, 0.5)
 
     def test_w_robust_search_space_univariate(self) -> None:
         # Check that if no transforms are needed, it is untouched.
         for multivariate in (True, False):
-            rss = get_robust_search_space(
-                multivariate=multivariate,
-                lb=self.target_lb,
-                ub=self.target_ub,
-            )
+            rss = get_robust_search_space(multivariate=multivariate, lb=0.0, ub=1.0)
             expected = str(rss)
-            t = self.transform_class(
-                search_space=rss,
-                observations=[],
-            )
+            t = UnitX(search_space=rss)
             self.assertEqual(expected, str(t.transform_search_space(rss)))
         # Error if distribution is multiplicative.
         rss = get_robust_search_space()
         rss.parameter_distributions[0].multiplicative = True
-        t = self.transform_class(
-            search_space=rss,
-            observations=[],
-        )
+        t = UnitX(search_space=rss)
         with self.assertRaisesRegex(NotImplementedError, "multiplicative"):
             t.transform_search_space(rss)
         # Correctly transform univariate additive distributions.
         rss = get_robust_search_space(lb=5.0, ub=10.0)
-        t = self.transform_class(
-            search_space=rss,
-            observations=[],
-        )
+        t = UnitX(search_space=rss)
         t.transform_search_space(rss)
         dists = rss.parameter_distributions
-        self.assertEqual(
-            dists[0].distribution_parameters["loc"], 0.2 * self.target_range
-        )
-        self.assertEqual(dists[0].distribution_parameters["scale"], self.target_range)
+        self.assertEqual(dists[0].distribution_parameters["loc"], 0.2)
+        self.assertEqual(dists[0].distribution_parameters["scale"], 1.0)
         self.assertEqual(dists[1].distribution_parameters["loc"], 0.0)
-        self.assertEqual(
-            dists[1].distribution_parameters["scale"], 0.2 * self.target_range
-        )
+        self.assertEqual(dists[1].distribution_parameters["scale"], 0.2)
         # Correctly transform environmental distributions.
         rss = get_robust_search_space(lb=5.0, ub=10.0)
         all_parameters = list(rss.parameters.values())
@@ -286,22 +247,19 @@ def test_w_robust_search_space_univariate(self) -> None:
             dist.distribution_parameters["loc"],
             t._normalize_value(1.0, (5.0, 10.0)),
         )
-        self.assertEqual(dist.distribution_parameters["scale"], self.target_range)
+        self.assertEqual(dist.distribution_parameters["scale"], 1.0)
         # Error if transform via loc / scale is not supported.
         rss = get_robust_search_space(use_discrete=True)
         rss.parameters["z"]._parameter_type = ParameterType.FLOAT
-        t = self.transform_class(
-            search_space=rss,
-            observations=[],
-        )
+        t = UnitX(search_space=rss)
         with self.assertRaisesRegex(UnsupportedError, "`loc` and `scale`"):
             t.transform_search_space(rss)
 
     def test_w_robust_search_space_multivariate(self) -> None:
         # Error if trying to transform non-normal multivariate distributions.
         rss = get_robust_search_space(multivariate=True)
         rss.parameter_distributions[0].distribution_class = "multivariate_t"
-        t = self.transform_class(
+        t = UnitX(
             search_space=rss,
             observations=[],
         )
@@ -310,25 +268,16 @@ def test_w_robust_search_space_multivariate(self) -> None:
         # Transform multivariate normal.
         rss = get_robust_search_space(multivariate=True)
         old_params = deepcopy(rss.parameter_distributions[0].distribution_parameters)
-        t = self.transform_class(
-            search_space=rss,
-            observations=[],
-        )
+        t = UnitX(search_space=rss)
         t.transform_search_space(rss)
         new_params = rss.parameter_distributions[0].distribution_parameters
         self.assertIsInstance(new_params["mean"], np.ndarray)
         self.assertIsInstance(new_params["cov"], np.ndarray)
         self.assertTrue(
-            np.allclose(
-                new_params["mean"],
-                np.asarray(old_params["mean"]) / 5.0 * self.target_range,
-            )
+            np.allclose(new_params["mean"], np.asarray(old_params["mean"]) / 5.0)
         )
         self.assertTrue(
-            np.allclose(
-                new_params["cov"],
-                np.asarray(old_params["cov"]) / ((5.0 / self.target_range) ** 2),
-            )
+            np.allclose(new_params["cov"], np.asarray(old_params["cov"]) / 25.0)
         )
         # Transform multivariate normal environmental distribution.
         rss = get_robust_search_space(multivariate=True)
@@ -339,18 +288,11 @@ def test_w_robust_search_space_multivariate(self) -> None:
             num_samples=rss.num_samples,
             environmental_variables=rss_params[:2],
         )
-        t = self.transform_class(
-            search_space=rss,
-            observations=[],
-        )
+        t = UnitX(search_space=rss)
         t.transform_search_space(rss)
         new_params = rss.parameter_distributions[0].distribution_parameters
         self.assertTrue(
-            np.allclose(
-                new_params["mean"],
-                np.asarray(old_params["mean"]) / 5.0 * self.target_range
-                + self.target_lb,
-            )
+            np.allclose(new_params["mean"], np.asarray(old_params["mean"]) / 5.0)
         )
         # Errors if mean / cov are of wrong shape.
         rss.parameter_distributions[0].distribution_parameters["mean"] = [1.0]
@@ -360,3 +302,42 @@ def test_w_robust_search_space_multivariate(self) -> None:
         rss.parameter_distributions[0].distribution_parameters["cov"] = [1.0]
         with self.assertRaisesRegex(UserInputError, "cov"):
             t.transform_search_space(rss)
+
+    def test_transform_experiment_data(self) -> None:
+        parameterizations = [
+            {"x": 1.0, "y": 1.5, "z": 1.0, "a": 1, "b": "b"},
+            {"x": 2.0, "y": 2.0, "z": 2.0, "a": 2, "b": "b"},
+        ]
+        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 `x` and `y` have been transformed.
+        expected = DataFrame(
+            index=transformed_data.arm_data.index,
+            data={
+                "x": [0.0, 0.5],
+                "y": [0.5, 1.0],
+            },
+            columns=["x", "y"],
+        )
+        assert_frame_equal(transformed_data.arm_data[["x", "y"]], expected)
+
+        # Remaining columns are unchanged.
+        # "z" is log-scale and "a" is in, so they're not transformed.
+        cols = ["z", "a", "b", "metadata"]
+        assert_frame_equal(
+            transformed_data.arm_data[cols], experiment_data.arm_data[cols]
+        )
+        # Observation data is unchanged.
+        assert_frame_equal(
+            transformed_data.observation_data, experiment_data.observation_data
+        )