Re-enable slow sensitivity plot tests (facebook#3893)

Summary:
Pull Request resolved: facebook#3893

These tests were removed in D74846187 because they were slow. D75569748 and D75712208 sped them up.

Reviewed By: saitcakmak

Differential Revision: D75712258

fbshipit-source-id: 24dd33697d9db4ec6bcc9ebf4a727181a8873906

Author: esantorella

ax/analysis/plotly/tests/test_sensitivity.py

@@ -5,6 +5,8 @@
 
 # pyre-strict
 
+from itertools import product
+
 from ax.adapter.registry import Generators
 from ax.analysis.analysis import (
     AnalysisBlobAnnotation,
@@ -20,39 +22,37 @@
 from ax.exceptions.core import UserInputError
 from ax.service.ax_client import AxClient, ObjectiveProperties
 from ax.utils.common.testutils import TestCase
-from ax.utils.testing.core_stubs import (
-    get_offline_experiments_subset,
-    get_online_experiments_subset,
-)
+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
 from pyre_extensions import assert_is_instance, none_throws
 
 
-class TestSensitivityAnalysisPlot(TestCase):
-    @mock_botorch_optimize
-    def setUp(self) -> None:
-        super().setUp()
-        self.client = AxClient()
-        self.client.create_experiment(
-            is_test=True,
-            name="foo",
-            parameters=[
-                {
-                    "name": "x",
-                    "type": "range",
-                    "bounds": [-1.0, 1.0],
-                }
-            ],
-            objectives={"bar": ObjectiveProperties(minimize=True)},
+@mock_botorch_optimize
+def get_test_client() -> AxClient:
+    client = AxClient()
+    client.create_experiment(
+        is_test=True,
+        name="foo",
+        parameters=[
+            {
+                "name": "x",
+                "type": "range",
+                "bounds": [-1.0, 1.0],
+            }
+        ],
+        objectives={"bar": ObjectiveProperties(minimize=True)},
+    )
+
+    for _ in range(10):
+        parameterization, trial_index = client.get_next_trial()
+        client.complete_trial(
+            trial_index=trial_index, raw_data={"bar": parameterization["x"] ** 2}
         )
+    return client
 
-        for _ in range(10):
-            parameterization, trial_index = self.client.get_next_trial()
-            self.client.complete_trial(
-                trial_index=trial_index, raw_data={"bar": parameterization["x"] ** 2}
-            )
 
+class TestSensitivityAnalysisPlot(TestCase):
     @mock_botorch_optimize
     def test_compute(self) -> None:
         client = Client()
@@ -120,10 +120,9 @@ def test_compute(self) -> None:
     @mock_botorch_optimize
     def test_compute_adhoc(self) -> None:
         metric_mapping = {"bar": "spunky"}
-        data = self.client.experiment.lookup_data()
-        adapter = Generators.BOTORCH_MODULAR(
-            experiment=self.client.experiment, data=data
-        )
+        client = get_test_client()
+        data = client.experiment.lookup_data()
+        adapter = Generators.BOTORCH_MODULAR(experiment=client.experiment, data=data)
         cards = compute_sensitivity_adhoc(adapter=adapter, labels=metric_mapping)
         self.assertEqual(len(cards), 1)
         card = cards[0]
@@ -133,24 +132,19 @@ def test_compute_adhoc(self) -> None:
     @mock_botorch_optimize
     @TestCase.ax_long_test(reason="Expensive to compute Sobol indicies")
     def test_online(self) -> None:
-        # Test SensitivityAnalysisPlot can be computed for a variety of experiments
-        # which resemble those we see in an online setting, in analogous tests we
-        # run all experiments with modifications to settings, however, this test
-        # is slow and so we limit the number of permutations we validate.
-        order = "total"  # most common
-
-        for experiment in get_online_experiments_subset():
-            for top_k in [None, 1]:
-                generation_strategy = get_default_generation_strategy_at_MBM_node(
-                    experiment=experiment
-                )
+        for experiment in get_online_experiments():
+            generation_strategy = get_default_generation_strategy_at_MBM_node(
+                experiment=experiment
+            )
+            # Select an arbitrary metric from the optimization config
+            metric_names = [
+                none_throws(experiment.optimization_config).objective.metric_names[0]
+            ]
+            for order, top_k in product(["first", "second", "total"], [None, 1]):
                 analysis = SensitivityAnalysisPlot(
-                    # Select and arbitrary metric from the optimization config
-                    metric_names=[
-                        none_throws(
-                            experiment.optimization_config
-                        ).objective.metric_names[0]
-                    ],
+                    metric_names=metric_names,
+                    # pyre-fixme: Incompatible parameter type [6]: It isn't sure
+                    # if "order" has one of the values specified by the Literal
                     order=order,
                     top_k=top_k,
                 )
@@ -162,28 +156,25 @@ def test_online(self) -> None:
     @mock_botorch_optimize
     @TestCase.ax_long_test(reason="Expensive to compute Sobol indicies")
     def test_offline(self) -> None:
-        # Test SensitivityAnalysisPlot can be computed for a variety of experiments
-        # which resemble those we see in an offline setting, in analogous tests we
-        # run all experiments with modifications to settings, however, this test
-        # is slow and so we limit the number of permutations we validate.
-        order = "total"  # most common
-
-        for experiment in get_offline_experiments_subset():
-            for top_k in [None, 1]:
-                generation_strategy = get_default_generation_strategy_at_MBM_node(
-                    experiment=experiment
-                )
+        for experiment in get_offline_experiments():
+            generation_strategy = get_default_generation_strategy_at_MBM_node(
+                experiment=experiment
+            )
+            # Select an arbitrary metric from the optimization config
+            metric_names = [
+                none_throws(experiment.optimization_config).objective.metric_names[0]
+            ]
+            for order, top_k in product(["first", "second", "total"], [None, 1]):
                 analysis = SensitivityAnalysisPlot(
-                    # Select and arbitrary metric from the optimization config
-                    metric_names=[
-                        none_throws(
-                            experiment.optimization_config
-                        ).objective.metric_names[0]
-                    ],
+                    metric_names=metric_names,
+                    # pyre-fixme: Incompatible parameter type [6]: It isn't sure
+                    # if "order" has one of the values specified by the Literal
                     order=order,
                     top_k=top_k,
                 )
 
+                # This prints a lot of warnings about y being constant
+                # because the first MOO experiment produces constant data
                 _ = analysis.compute(
                     experiment=experiment, generation_strategy=generation_strategy
                 )