Consolidate ax/generator tests

Summary: Part of a 19-diff stack to consolidate repetitive tests across Ax using `subTest`. Consolidate 8 test files in ax/generators/ — adds subTest to Thompson sampler weight configs, kernel tests, covariance module argparse, and generator utils tests.

Differential Revision: D95604917

Author: eonofrey

ax/generators/tests/test_discrete.py

@@ -16,11 +16,6 @@ def test_discrete_model_get_state(self) -> None:
         discrete_model = DiscreteGenerator()
         self.assertEqual(discrete_model._get_state(), {})
 
-    def test_discrete_model_feature_importances(self) -> None:
-        discrete_model = DiscreteGenerator()
-        with self.assertRaises(NotImplementedError):
-            discrete_model.feature_importances()
-
     def test_DiscreteGeneratorFit(self) -> None:
         discrete_model = DiscreteGenerator()
         discrete_model.fit(
@@ -38,16 +33,18 @@ def test_discreteModelPredict(self) -> None:
         with self.assertRaises(NotImplementedError):
             discrete_model.predict([[0]])
 
-    def test_discreteModelGen(self) -> None:
+    def test_not_implemented_methods(self) -> None:
         discrete_model = DiscreteGenerator()
-        with self.assertRaises(NotImplementedError):
-            discrete_model.gen(
+        cases = {
+            "feature_importances": lambda: discrete_model.feature_importances(),
+            "gen": lambda: discrete_model.gen(
                 n=1, parameter_values=[[0, 1]], objective_weights=np.array([1])
-            )
-
-    def test_discreteModelCrossValidate(self) -> None:
-        discrete_model = DiscreteGenerator()
-        with self.assertRaises(NotImplementedError):
-            discrete_model.cross_validate(
+            ),
+            "cross_validate": lambda: discrete_model.cross_validate(
                 Xs_train=[[[0]]], Ys_train=[[1]], Yvars_train=[[1]], X_test=[[1]]
-            )
+            ),
+        }
+        for method_name, call in cases.items():
+            with self.subTest(method=method_name):
+                with self.assertRaises(NotImplementedError):
+                    call()

ax/generators/tests/test_random.py

@@ -31,21 +31,24 @@ def test_seed(self) -> None:
 
     def test_state(self) -> None:
         for model in (self.random_model, RandomGenerator(seed=5)):
-            state = model._get_state()
-            self.assertEqual(state["seed"], model.seed)
-            self.assertEqual(state["init_position"], model.init_position)
+            with self.subTest(seed=model.seed):
+                state = model._get_state()
+                self.assertEqual(state["seed"], model.seed)
+                self.assertEqual(state["init_position"], model.init_position)
 
-    def test_RandomGeneratorGenSamples(self) -> None:
-        with self.assertRaises(NotImplementedError):
-            self.random_model._gen_samples(
+    def test_not_implemented_methods(self) -> None:
+        cases = {
+            "_gen_samples": lambda: self.random_model._gen_samples(
                 n=1, tunable_d=1, bounds=np.array([[0.0, 1.0]])
-            )
-
-    def test_RandomGeneratorGenUnconstrained(self) -> None:
-        with self.assertRaises(NotImplementedError):
-            self.random_model._gen_unconstrained(
+            ),
+            "_gen_unconstrained": lambda: self.random_model._gen_unconstrained(
                 n=1, d=2, tunable_feature_indices=np.array([], dtype=int)
-            )
+            ),
+        }
+        for method_name, call in cases.items():
+            with self.subTest(method=method_name):
+                with self.assertRaises(NotImplementedError):
+                    call()
 
     def test_ConvertEqualityConstraints(self) -> None:
         fixed_features = {3: 0.7, 1: 0.5}

ax/generators/tests/test_thompson.py

@@ -57,41 +57,80 @@ def test_ThompsonSampler(self) -> None:
         self.assertEqual(len(gen_metadata["arms_to_weights"]), 4)
         self.assertEqual(gen_metadata["best_x"], arms[0])
 
+    def test_ThompsonSamplerWeightConfigs(self) -> None:
+        for label, min_weight, uniform_weights, expected_arms, expected_weights in [
+            (
+                "min_weight=0.01",
+                0.01,
+                False,
+                [[4, 4], [3, 3], [2, 2]],
+                [3 * i for i in [0.725, 0.225, 0.05]],
+            ),
+            (
+                "uniform_weights",
+                0.0,
+                True,
+                [[4, 4], [3, 3], [2, 2]],
+                [1.0, 1.0, 1.0],
+            ),
+        ]:
+            with self.subTest(config=label):
+                np.random.seed(0)
+                generator = ThompsonSampler(
+                    min_weight=min_weight, uniform_weights=uniform_weights
+                )
+                generator.fit(
+                    Xs=self.Xs,
+                    Ys=self.Ys,
+                    Yvars=self.Yvars,
+                    parameter_values=self.parameter_values,
+                    outcome_names=self.outcome_names,
+                )
+                arms, weights, _ = generator.gen(
+                    n=3,
+                    parameter_values=self.parameter_values,
+                    objective_weights=np.ones(1),
+                )
+                self.assertEqual(arms, expected_arms)
+                for weight, expected_weight in zip(weights, expected_weights):
+                    self.assertAlmostEqual(weight, expected_weight, 1)
+
     def test_ThompsonSamplerValidation(self) -> None:
         generator = ThompsonSampler(min_weight=0.01)
 
-        # all Xs are not the same
-        with self.assertRaises(ValueError):
-            generator.fit(
-                Xs=[[[1, 1], [2, 2], [3, 3], [4, 4]], [[1, 1], [2, 2], [4, 4]]],
-                Ys=self.Ys,
-                Yvars=self.Yvars,
-                parameter_values=self.parameter_values,
-                outcome_names=self.outcome_names,
-            )
+        with self.subTest(case="mismatched_Xs"):
+            with self.assertRaises(ValueError):
+                generator.fit(
+                    Xs=[[[1, 1], [2, 2], [3, 3], [4, 4]], [[1, 1], [2, 2], [4, 4]]],
+                    Ys=self.Ys,
+                    Yvars=self.Yvars,
+                    parameter_values=self.parameter_values,
+                    outcome_names=self.outcome_names,
+                )
 
-        # multiple observations per parameterization
-        with self.assertRaises(ValueError):
+        with self.subTest(case="duplicate_parameterizations"):
+            with self.assertRaises(ValueError):
+                generator.fit(
+                    Xs=[[[1, 1], [2, 2], [2, 2]]],
+                    Ys=self.Ys,
+                    Yvars=self.Yvars,
+                    parameter_values=self.parameter_values,
+                    outcome_names=self.outcome_names,
+                )
+
+        with self.subTest(case="similar_but_different_observations"):
+            # these are not the same observations, so should not error
             generator.fit(
-                Xs=[[[1, 1], [2, 2], [2, 2]]],
+                Xs=[[[1, 1], [2.0, 2], [2, 2]]],
                 Ys=self.Ys,
                 Yvars=self.Yvars,
                 parameter_values=self.parameter_values,
                 outcome_names=self.outcome_names,
             )
 
-        # these are not the same observations, so should not error
-        generator.fit(
-            Xs=[[[1, 1], [2.0, 2], [2, 2]]],
-            Ys=self.Ys,
-            Yvars=self.Yvars,
-            parameter_values=self.parameter_values,
-            outcome_names=self.outcome_names,
-        )
-
-        # requires objective weights
-        with self.assertRaises(ValueError):
-            generator.gen(5, self.parameter_values, objective_weights=None)
+        with self.subTest(case="missing_objective_weights"):
+            with self.assertRaises(ValueError):
+                generator.gen(5, self.parameter_values, objective_weights=None)
 
     def test_ThompsonSamplerTopKError(self) -> None:
         generator = ThompsonSampler(topk=5)
@@ -156,45 +195,6 @@ def test_TopTwo_alters_weights_vs_TopOne(self) -> None:
         # 4) Monotonicity in the final TTTS distribution still holds
         self.assertTrue(full_w2[3] > full_w2[2] > full_w2[1] > full_w2[0])
 
-    def test_ThompsonSamplerMinWeight(self) -> None:
-        np.random.seed(0)
-        generator = ThompsonSampler(min_weight=0.01)
-        generator.fit(
-            Xs=self.Xs,
-            Ys=self.Ys,
-            Yvars=self.Yvars,
-            parameter_values=self.parameter_values,
-            outcome_names=self.outcome_names,
-        )
-        arms, weights, _ = generator.gen(
-            n=3,
-            parameter_values=self.parameter_values,
-            objective_weights=np.ones(1),
-        )
-        self.assertEqual(arms, [[4, 4], [3, 3], [2, 2]])
-        for weight, expected_weight in zip(
-            weights, [3 * i for i in [0.725, 0.225, 0.05]]
-        ):
-            self.assertAlmostEqual(weight, expected_weight, 1)
-
-    def test_ThompsonSamplerUniformWeights(self) -> None:
-        generator = ThompsonSampler(min_weight=0.0, uniform_weights=True)
-        generator.fit(
-            Xs=self.Xs,
-            Ys=self.Ys,
-            Yvars=self.Yvars,
-            parameter_values=self.parameter_values,
-            outcome_names=self.outcome_names,
-        )
-        arms, weights, _ = generator.gen(
-            n=3,
-            parameter_values=self.parameter_values,
-            objective_weights=np.ones(1),
-        )
-        self.assertEqual(arms, [[4, 4], [3, 3], [2, 2]])
-        for weight, expected_weight in zip(weights, [1.0, 1.0, 1.0]):
-            self.assertAlmostEqual(weight, expected_weight, 1)
-
     def test_ThompsonSamplerInfeasible(self) -> None:
         generator = ThompsonSampler(min_weight=0.9)
         generator.fit(
@@ -302,9 +302,12 @@ def test_ThompsonSamplerNonPositiveN(self) -> None:
             outcome_names=self.outcome_names,
         )
         for n in (-1, 0):
-            with self.assertRaisesRegex(ValueError, "ThompsonSampler requires n > 0"):
-                generator.gen(
-                    n=n,
-                    parameter_values=self.parameter_values,
-                    objective_weights=np.ones(1),
-                )
+            with self.subTest(n=n):
+                with self.assertRaisesRegex(
+                    ValueError, "ThompsonSampler requires n > 0"
+                ):
+                    generator.gen(
+                        n=n,
+                        parameter_values=self.parameter_values,
+                        objective_weights=np.ones(1),
+                    )

ax/generators/tests/test_torch.py

@@ -39,19 +39,25 @@ def test_TorchModelFit(self) -> None:
             ),
         )
 
-    def test_TorchModelPredict(self) -> None:
+    def test_not_implemented_methods(self) -> None:
         torch_model = TorchGenerator()
-        with self.assertRaises(NotImplementedError):
-            torch_model.predict(torch.zeros(1))
-
-    def test_TorchModelGen(self) -> None:
-        torch_model = TorchGenerator()
-        with self.assertRaises(NotImplementedError):
-            torch_model.gen(
+        cases = {
+            "predict": lambda: torch_model.predict(torch.zeros(1)),
+            "gen": lambda: torch_model.gen(
                 n=1,
                 search_space_digest=self.search_space_digest,
                 torch_opt_config=self.torch_opt_config,
-            )
+            ),
+            "cross_validate": lambda: torch_model.cross_validate(
+                datasets=[self.dataset],
+                X_test=torch.ones(1),
+                search_space_digest=SearchSpaceDigest(feature_names=[], bounds=[]),
+            ),
+        }
+        for method_name, call in cases.items():
+            with self.subTest(method=method_name):
+                with self.assertRaises(NotImplementedError):
+                    call()
 
     def test_NumpyTorchBestPoint(self) -> None:
         torch_model = TorchGenerator()
@@ -60,12 +66,3 @@ def test_NumpyTorchBestPoint(self) -> None:
             torch_opt_config=self.torch_opt_config,
         )
         self.assertIsNone(x)
-
-    def test_TorchModelCrossValidate(self) -> None:
-        torch_model = TorchGenerator()
-        with self.assertRaises(NotImplementedError):
-            torch_model.cross_validate(
-                datasets=[self.dataset],
-                X_test=torch.ones(1),
-                search_space_digest=SearchSpaceDigest(feature_names=[], bounds=[]),
-            )