Move relativize functions out of statstools to avoid circular dependency in imports (facebook#4341)

Summary:

To resolve a circular dependency between `ax.core` and `statstools`, this diff moves the `relativize`, `unrelativize`, and `relativize_data` functions out of statstools. The `Data` class now provides a relativized method to replace relativize_data, and all relevant imports, dependencies, and unit tests have been updated accordingly.

Differential Revision: D83095707

Author: shrutipatel31

ax/adapter/transforms/relativize.py

@@ -27,7 +27,7 @@
 from ax.core.search_space import SearchSpace
 from ax.exceptions.core import DataRequiredError
 from ax.generators.types import TConfig
-from ax.utils.stats.statstools import relativize, unrelativize
+from ax.utils.stats.math_utils import relativize, unrelativize
 from pyre_extensions import none_throws
 
 if TYPE_CHECKING:

ax/adapter/transforms/tests/test_relativize_transform.py

@@ -36,7 +36,6 @@
 from ax.generators.base import Generator
 from ax.metrics.branin import BraninMetric
 from ax.utils.common.testutils import TestCase
-from ax.utils.stats.statstools import relativize_data
 from ax.utils.testing.core_stubs import (
     get_branin_data_batch,
     get_branin_experiment,
@@ -352,8 +351,7 @@ def test_multitask_data(self) -> None:
         )
         relative_observations = observations_from_data(
             experiment=experiment,
-            data=relativize_data(
-                data=data,
+            data=data.relativize(
                 status_quo_name="status_quo",
                 as_percent=True,
                 include_sq=True,
@@ -387,7 +385,7 @@ def test_multitask_data(self) -> None:
         )
 
         # not checking RelativizeWithConstantControl here
-        # because relativize_data uses delta method
+        # because data.relativize uses delta method
         transform = Relativize(search_space=None, adapter=adapter)
 
         relative_obs_t = transform.transform_observations(observations)

ax/adapter/transforms/tests/test_transform_to_new_sq.py

@@ -22,7 +22,7 @@
 from ax.exceptions.core import DataRequiredError
 from ax.generators.base import Generator
 from ax.utils.common.testutils import TestCase
-from ax.utils.stats.statstools import relativize
+from ax.utils.stats.math_utils import relativize
 from ax.utils.testing.core_stubs import (
     get_branin_data_batch,
     get_branin_experiment,

ax/adapter/transforms/transform_to_new_sq.py

@@ -23,7 +23,7 @@
 from ax.core.utils import get_target_trial_index
 from ax.generators.types import TConfig
 from ax.utils.common.logger import get_logger
-from ax.utils.stats.statstools import relativize, unrelativize
+from ax.utils.stats.math_utils import relativize, unrelativize
 from pyre_extensions import assert_is_instance, none_throws
 
 if TYPE_CHECKING:

ax/analysis/tests/test_utils.py

@@ -20,7 +20,6 @@
 from ax.core.trial_status import TrialStatus  # noqa
 from ax.exceptions.core import UserInputError
 from ax.utils.common.testutils import TestCase
-from ax.utils.stats.statstools import relativize_data
 from ax.utils.testing.core_stubs import get_offline_experiments, get_online_experiments
 from ax.utils.testing.mock import mock_botorch_optimize
 from ax.utils.testing.modeling_stubs import get_default_generation_strategy_at_MBM_node
@@ -676,7 +675,7 @@ def test_online(self) -> None:
         # resemble those we see in an online setting.
         for experiment in get_online_experiments():
             data = experiment.lookup_data()
-            rel_df = relativize_data(data).df
+            rel_df = data.relativize().df
             raw_df = data.df
             metric_name = next(iter(experiment.metrics.keys()))
             raw_arm_value = raw_df[

ax/analysis/utils.py

@@ -27,7 +27,7 @@
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.utils.common.constants import Keys
 from ax.utils.common.logger import get_logger
-from ax.utils.stats.statstools import relativize
+from ax.utils.stats.math_utils import relativize
 from botorch.utils.probability.utils import compute_log_prob_feas_from_bounds
 from pyre_extensions import none_throws
 

ax/core/data.py

@@ -27,6 +27,7 @@
     TClassDecoderRegistry,
     TDecoderRegistry,
 )
+from ax.utils.stats.math_utils import relativize
 from pyre_extensions import assert_is_instance
 
 logger: Logger = get_logger(__name__)
@@ -392,6 +393,91 @@ def from_multiple_data(cls, data: Iterable[Data]) -> Data:
         """
         return cls.from_multiple(data=data)
 
+    def relativize(
+        self,
+        status_quo_name: str = "status_quo",
+        as_percent: bool = False,
+        include_sq: bool = False,
+        bias_correction: bool = True,
+        control_as_constant: bool = False,
+    ) -> "Data":
+        """Relativize this data object w.r.t. a status_quo arm.
+
+        Args:
+            status_quo_name: The name of the status_quo arm.
+            as_percent: If True, return results as percentage change.
+            include_sq: Include status quo in final df.
+            bias_correction: Whether to apply bias correction when computing relativized
+                metric values. Uses a second-order Taylor expansion for approximating
+                the means and standard errors or the ratios, see
+                ax.utils.stats.statstools.relativize for more details.
+            control_as_constant: If true, control is treated as a constant.
+                bias_correction is ignored when this is true.
+
+        Returns:
+            The new data object with the relativized metrics (excluding the
+                status_quo arm)
+
+        """
+
+        df = self.df.copy()
+        grp_cols = list(
+            {"trial_index", "metric_name", "random_split"}.intersection(
+                df.columns.values
+            )
+        )
+
+        grouped_df = df.groupby(grp_cols)
+        dfs = []
+        for grp in grouped_df.groups.keys():
+            subgroup_df = grouped_df.get_group(grp)
+            is_sq = subgroup_df["arm_name"] == status_quo_name
+
+            # Check if status quo exists in this subgroup
+            sq_data = (
+                subgroup_df[is_sq][["mean", "sem"]].drop_duplicates().values.flatten()
+            )
+            if len(sq_data) == 0:
+                # No status quo in this subgroup, skip relativization
+                continue
+            elif len(sq_data) != 2:
+                raise ValueError(
+                    f"Expected exactly 2 values (mean, sem) for status quo, "
+                    f"got {len(sq_data)}"
+                )
+
+            sq_mean, sq_sem = sq_data
+
+            # rm status quo from final df to relativize
+            if not include_sq:
+                subgroup_df = subgroup_df[~is_sq]
+            means_rel, sems_rel = relativize(
+                means_t=subgroup_df["mean"].values,
+                sems_t=subgroup_df["sem"].values,
+                mean_c=sq_mean,
+                sem_c=sq_sem,
+                as_percent=as_percent,
+                bias_correction=bias_correction,
+                control_as_constant=control_as_constant,
+            )
+            dfs.append(
+                pd.concat(
+                    [
+                        subgroup_df.drop(["mean", "sem"], axis=1),
+                        pd.DataFrame(
+                            np.array([means_rel, sems_rel]).T,
+                            columns=["mean", "sem"],
+                            index=subgroup_df.index,
+                        ),
+                    ],
+                    axis=1,
+                )
+            )
+        df_rel = pd.concat(dfs, axis=0)
+        if include_sq:
+            df_rel.loc[df_rel["arm_name"] == status_quo_name, "sem"] = 0.0
+        return Data(df_rel)
+
     def clone(self) -> Data:
         """Returns a new Data object with the same underlying dataframe."""
         return Data(df=deepcopy(self.df))

ax/core/tests/test_data.py

@@ -8,6 +8,8 @@
 
 from unittest.mock import patch
 
+import numpy as np
+
 import pandas as pd
 from ax.core.data import Data
 from ax.core.map_data import MAP_KEY, MapData
@@ -366,3 +368,141 @@ def test_filter(self) -> None:
         filtered = data.filter(metric_names=["a"])
         self.assertEqual(len(filtered.df), 3)
         self.assertEqual(set(filtered.df["metric_name"]), {"a"})
+
+
+class RelativizeDataTest(TestCase):
+    def setUp(self) -> None:
+        self.df = pd.DataFrame(
+            [
+                {
+                    "trial_index": 0,
+                    "mean": 2,
+                    "sem": 0,
+                    "metric_name": "foobar",
+                    "metric_signature": "foobar",
+                    "arm_name": "status_quo",
+                },
+                {
+                    "trial_index": 0,
+                    "mean": 5,
+                    "sem": 0,
+                    "metric_name": "foobaz",
+                    "metric_signature": "foobaz",
+                    "arm_name": "status_quo",
+                },
+                {
+                    "trial_index": 0,
+                    "mean": 1,
+                    "sem": 0,
+                    "metric_name": "foobar",
+                    "metric_signature": "foobar",
+                    "arm_name": "0_0",
+                },
+                {
+                    "trial_index": 0,
+                    "mean": 10,
+                    "sem": 0,
+                    "metric_name": "foobaz",
+                    "metric_signature": "foobaz",
+                    "arm_name": "0_0",
+                },
+            ]
+        )
+
+        self.expected_relativized_df = pd.DataFrame(
+            [
+                {
+                    "trial_index": 0,
+                    "mean": -0.5,
+                    "sem": 0,
+                    "metric_name": "foobar",
+                    "metric_signature": "foobar",
+                    "arm_name": "0_0",
+                },
+                {
+                    "trial_index": 0,
+                    "mean": 1,
+                    "sem": 0,
+                    "metric_name": "foobaz",
+                    "metric_signature": "foobaz",
+                    "arm_name": "0_0",
+                },
+            ]
+        )
+        self.expected_relativized_df_with_sq = pd.DataFrame(
+            [
+                {
+                    "trial_index": 0,
+                    "mean": 0,
+                    "sem": 0,
+                    "metric_name": "foobar",
+                    "metric_signature": "foobar",
+                    "arm_name": "status_quo",
+                },
+                {
+                    "trial_index": 0,
+                    "mean": -0.5,
+                    "sem": 0,
+                    "metric_name": "foobar",
+                    "metric_signature": "foobar",
+                    "arm_name": "0_0",
+                },
+                {
+                    "trial_index": 0,
+                    "mean": 0,
+                    "sem": 0,
+                    "metric_name": "foobaz",
+                    "metric_signature": "foobaz",
+                    "arm_name": "status_quo",
+                },
+                {
+                    "trial_index": 0,
+                    "mean": 1,
+                    "sem": 0,
+                    "metric_name": "foobaz",
+                    "metric_signature": "foobaz",
+                    "arm_name": "0_0",
+                },
+            ]
+        )
+
+    def test_relativize_data(self) -> None:
+        data = Data(
+            df=self.df,
+        )
+        expected_relativized_data = Data(df=self.expected_relativized_df)
+
+        expected_relativized_data_with_sq = Data(
+            df=self.expected_relativized_df_with_sq
+        )
+
+        actual_relativized_data = data.relativize()
+        self.assertEqual(expected_relativized_data, actual_relativized_data)
+
+        actual_relativized_data_with_sq = data.relativize(include_sq=True)
+        self.assertEqual(
+            expected_relativized_data_with_sq, actual_relativized_data_with_sq
+        )
+
+    def test_relativize_data_no_sem(self) -> None:
+        df = self.df.copy()
+        df["sem"] = np.nan
+        data = Data(df=df)
+
+        expected_relativized_df = self.expected_relativized_df.copy()
+        expected_relativized_df["sem"] = np.nan
+        expected_relativized_data = Data(df=expected_relativized_df)
+
+        expected_relativized_df_with_sq = self.expected_relativized_df_with_sq.copy()
+        expected_relativized_df_with_sq.loc[
+            expected_relativized_df_with_sq["arm_name"] != "status_quo", "sem"
+        ] = np.nan
+        expected_relativized_data_with_sq = Data(df=expected_relativized_df_with_sq)
+
+        actual_relativized_data = data.relativize()
+        self.assertEqual(expected_relativized_data, actual_relativized_data)
+
+        actual_relativized_data_with_sq = data.relativize(include_sq=True)
+        self.assertEqual(
+            expected_relativized_data_with_sq, actual_relativized_data_with_sq
+        )

ax/plot/pareto_utils.py

@@ -41,7 +41,7 @@
 from ax.exceptions.core import AxError, UnsupportedError, UserInputError
 from ax.generators.torch_base import TorchGenerator
 from ax.utils.common.logger import get_logger
-from ax.utils.stats.statstools import relativize
+from ax.utils.stats.math_utils import relativize
 from botorch.acquisition.monte_carlo import qSimpleRegret
 from botorch.utils.multi_objective import is_non_dominated
 from botorch.utils.multi_objective.hypervolume import infer_reference_point

ax/plot/scatter.py

@@ -45,7 +45,7 @@
 )
 from ax.utils.common.logger import get_logger
 from ax.utils.common.typeutils import assert_is_instance_optional
-from ax.utils.stats.statstools import relativize
+from ax.utils.stats.math_utils import relativize
 from plotly import subplots
 
 logger: Logger = get_logger(__name__)