Stop some warnings in unit tests (#1992)

Summary:
## Motivation

Warning output from unit tests sometimes indicates a serious problem and sometimes merely clutters the output so we can't notice the serious problems. These warnings are the latter:
*  InputDataWarning: Input data is not contained to the unit cube. Please consider min-max scaling the input data (occurred 194 times, now 0)
* BadInitialCandidatesWarning: Unable to find non-zero acquisition function values - initial conditions are being selected randomly. (occurred 40 times, now 0)
*  The first positional argument of samplers, `num_samples`, has been deprecated and replaced with `sample_shape`, which expects a `torch.Size` object.' (occurred 35 times, now 0)

### Have you read the [Contributing Guidelines on pull requests](https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Yes

Pull Request resolved: #1992

Test Plan:
Units

## Related PRs

#1792, #1539

Reviewed By: saitcakmak

Differential Revision: D48530764

Pulled By: esantorella

fbshipit-source-id: c5b1898ce8156a6f02550acb09bd5eba0c157c5e

botorch/utils/testing.py

@@ -43,25 +43,31 @@ class BotorchTestCase(TestCase):
 
     device = torch.device("cpu")
 
-    def setUp(self):
+    def setUp(self, suppress_input_warnings: bool = True) -> None:
         warnings.resetwarnings()
         settings.debug._set_state(False)
         warnings.simplefilter("always", append=True)
-        warnings.filterwarnings(
-            "ignore",
-            message="The model inputs are of type",
-            category=UserWarning,
-        )
-        warnings.filterwarnings(
-            "ignore",
-            message="Non-strict enforcement of botorch tensor conventions.",
-            category=BotorchTensorDimensionWarning,
-        )
-        warnings.filterwarnings(
-            "ignore",
-            message="Input data is not standardized.",
-            category=InputDataWarning,
-        )
+        if suppress_input_warnings:
+            warnings.filterwarnings(
+                "ignore",
+                message="The model inputs are of type",
+                category=UserWarning,
+            )
+            warnings.filterwarnings(
+                "ignore",
+                message="Non-strict enforcement of botorch tensor conventions.",
+                category=BotorchTensorDimensionWarning,
+            )
+            warnings.filterwarnings(
+                "ignore",
+                message="Input data is not standardized.",
+                category=InputDataWarning,
+            )
+            warnings.filterwarnings(
+                "ignore",
+                message="Input data is not contained to the unit cube.",
+                category=InputDataWarning,
+            )
 
     def assertAllClose(
         self,

test/acquisition/test_input_constructors.py

@@ -97,6 +97,7 @@ class DummyAcquisitionFunction(AcquisitionFunction):
 
 class InputConstructorBaseTestCase:
     def setUp(self) -> None:
+        super().setUp()
         self.mock_model = MockModel(
             posterior=MockPosterior(mean=None, variance=None, base_shape=(1,))
         )

test/acquisition/test_monte_carlo.py

@@ -556,7 +556,7 @@ def test_cache_root(self):
             "prune_baseline": False,
             "cache_root": True,
             "posterior_transform": ScalarizedPosteriorTransform(weights=torch.ones(m)),
-            "sampler": SobolQMCNormalSampler(5),
+            "sampler": SobolQMCNormalSampler(sample_shape=torch.Size([5])),
         }
         acqf = qNoisyExpectedImprovement(**nei_args)
         X = torch.randn_like(X_baseline)

test/acquisition/test_preference.py

@@ -18,7 +18,8 @@
 
 
 class TestPreferenceAcquisitionFunctions(BotorchTestCase):
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         self.twargs = {"dtype": torch.double}
         self.X_dim = 3
         self.Y_dim = 2

test/acquisition/test_prior_guided.py

@@ -39,6 +39,7 @@ def get_weighted_val(ei_val, prob, exponent, use_log):
 
 class TestPriorGuidedAcquisitionFunction(BotorchTestCase):
     def setUp(self):
+        super().setUp()
         self.prior = DummyPrior()
         self.train_X = torch.rand(5, 3, dtype=torch.double, device=self.device)
         self.train_Y = self.train_X.norm(dim=-1, keepdim=True)

test/models/test_gpytorch.py

@@ -296,7 +296,7 @@ def test_fantasize_flag(self):
         self.assertFalse(model.last_fantasize_flag)
         model.posterior(test_X)
         self.assertFalse(model.last_fantasize_flag)
-        model.fantasize(test_X, SobolQMCNormalSampler(2))
+        model.fantasize(test_X, SobolQMCNormalSampler(sample_shape=torch.Size([2])))
         self.assertTrue(model.last_fantasize_flag)
         model.last_fantasize_flag = False
         with fantasize():

test/models/test_model_list_gp_regression.py

@@ -382,7 +382,9 @@ def test_fantasize(self):
         m1 = SingleTaskGP(torch.rand(5, 2), torch.rand(5, 1)).eval()
         m2 = SingleTaskGP(torch.rand(5, 2), torch.rand(5, 1)).eval()
         modellist = ModelListGP(m1, m2)
-        fm = modellist.fantasize(torch.rand(3, 2), sampler=IIDNormalSampler(2))
+        fm = modellist.fantasize(
+            torch.rand(3, 2), sampler=IIDNormalSampler(sample_shape=torch.Size([2]))
+        )
         self.assertIsInstance(fm, ModelListGP)
         for i in range(2):
             fm_i = fm.models[i]
@@ -391,8 +393,8 @@ def test_fantasize(self):
             self.assertEqual(fm_i.train_targets.shape, torch.Size([2, 8]))
 
         # test decoupled
-        sampler1 = IIDNormalSampler(2)
-        sampler2 = IIDNormalSampler(2)
+        sampler1 = IIDNormalSampler(sample_shape=torch.Size([2]))
+        sampler2 = IIDNormalSampler(sample_shape=torch.Size([2]))
         eval_mask = torch.tensor(
             [[1, 0], [0, 1], [1, 0]],
             dtype=torch.bool,
@@ -457,7 +459,7 @@ def _get_fant_mean(
                     return fant.posterior(target_x).mean.mean(dim=(-2, -3))
 
                 # ~0
-                sampler = IIDNormalSampler(10, seed=0)
+                sampler = IIDNormalSampler(sample_shape=torch.Size([10]), seed=0)
                 fant_mean_with_manual_transform = _get_fant_mean(
                     model_manually_transformed, sampler=sampler
                 )
@@ -490,8 +492,8 @@ def _get_fant_mean(
                 )
                 # test decoupled
                 sampler = ListSampler(
-                    IIDNormalSampler(10, seed=0),
-                    IIDNormalSampler(10, seed=0),
+                    IIDNormalSampler(sample_shape=torch.Size([10]), seed=0),
+                    IIDNormalSampler(sample_shape=torch.Size([10]), seed=0),
                 )
                 fant_mean_with_manual_transform = _get_fant_mean(
                     model_manually_transformed,
@@ -539,7 +541,7 @@ def test_fantasize_with_outcome_transform_fixed_noise(self) -> None:
         100 at x=0. If transforms are not properly applied, we'll get answers
         on the order of ~1. Answers between 99 and 101 are acceptable.
         """
-        n_fants = 20
+        n_fants = torch.Size([20])
         y_at_low_x = 100.0
         y_at_high_x = -40.0
 

test/models/utils/test_assorted.py

@@ -72,6 +72,11 @@ def test_add_output_dim(self):
 
 
 class TestInputDataChecks(BotorchTestCase):
+    def setUp(self) -> None:
+        # The super class usually disables input data warnings in unit tests.
+        # Don't do that here.
+        super().setUp(suppress_input_warnings=False)
+
     def test_check_no_nans(self):
         check_no_nans(torch.tensor([1.0, 2.0]))
         with self.assertRaises(InputDataError):
@@ -87,12 +92,10 @@ def test_check_min_max_scaling(self):
                     any(issubclass(w.category, InputDataWarning) for w in ws)
                 )
             check_min_max_scaling(X=X, raise_on_fail=True)
-            with warnings.catch_warnings(record=True) as ws:
+            with self.assertWarnsRegex(
+                expected_warning=InputDataWarning, expected_regex="not scaled"
+            ):
                 check_min_max_scaling(X=X, strict=True)
-                self.assertTrue(
-                    any(issubclass(w.category, InputDataWarning) for w in ws)
-                )
-                self.assertTrue(any("not scaled" in str(w.message) for w in ws))
             with self.assertRaises(InputDataError):
                 check_min_max_scaling(X=X, strict=True, raise_on_fail=True)
             # check proper input

test/optim/test_fit.py

@@ -25,7 +25,8 @@
 
 
 class TestFitGPyTorchMLLScipy(BotorchTestCase):
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         self.mlls = {}
         with torch.random.fork_rng():
             torch.manual_seed(0)
@@ -172,7 +173,8 @@ def _assert_np_array_is_float64_type(array) -> bool:
 
 
 class TestFitGPyTorchMLLTorch(BotorchTestCase):
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         self.mlls = {}
         with torch.random.fork_rng():
             torch.manual_seed(0)
@@ -236,7 +238,8 @@ def _test_fit_gpytorch_mll_torch(self, mll):
 
 
 class TestFitGPyTorchScipy(BotorchTestCase):
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         self.mlls = {}
         with torch.random.fork_rng():
             torch.manual_seed(0)
@@ -372,6 +375,7 @@ def _test_fit_gpytorch_scipy(self, mll):
 
 class TestFitGPyTorchTorch(BotorchTestCase):
     def setUp(self):
+        super().setUp()
         self.mlls = {}
         with torch.random.fork_rng():
             torch.manual_seed(0)

test/optim/test_initializers.py

@@ -188,7 +188,10 @@ def test_gen_batch_initial_conditions(self):
                     MockAcquisitionFunction,
                     "__call__",
                     wraps=mock_acqf.__call__,
-                ) as mock_acqf_call:
+                ) as mock_acqf_call, warnings.catch_warnings():
+                    warnings.simplefilter(
+                        "ignore", category=BadInitialCandidatesWarning
+                    )
                     batch_initial_conditions = gen_batch_initial_conditions(
                         acq_function=mock_acqf,
                         bounds=bounds,
@@ -248,6 +251,9 @@ def test_gen_batch_initial_conditions_highdim(self):
                 [True, False], [None, 1234], [None, ffs_map], [True, False]
             ):
                 with warnings.catch_warnings(record=True) as ws, settings.debug(True):
+                    warnings.simplefilter(
+                        "ignore", category=BadInitialCandidatesWarning
+                    )
                     batch_initial_conditions = gen_batch_initial_conditions(
                         acq_function=MockAcquisitionFunction(),
                         bounds=bounds,
@@ -279,19 +285,17 @@ def test_gen_batch_initial_conditions_highdim(self):
                             torch.all(batch_initial_conditions[..., idx] == val)
                         )
 
-    def test_gen_batch_initial_conditions_warning(self):
+    def test_gen_batch_initial_conditions_warning(self) -> None:
         for dtype in (torch.float, torch.double):
             bounds = torch.tensor([[0, 0], [1, 1]], device=self.device, dtype=dtype)
             samples = torch.zeros(10, 1, 2, device=self.device, dtype=dtype)
-            with ExitStack() as es:
-                ws = es.enter_context(warnings.catch_warnings(record=True))
-                es.enter_context(settings.debug(True))
-                es.enter_context(
-                    mock.patch(
-                        "botorch.optim.initializers.draw_sobol_samples",
-                        return_value=samples,
-                    )
-                )
+            with self.assertWarnsRegex(
+                expected_warning=BadInitialCandidatesWarning,
+                expected_regex="Unable to find non-zero acquisition",
+            ), mock.patch(
+                "botorch.optim.initializers.draw_sobol_samples",
+                return_value=samples,
+            ):
                 batch_initial_conditions = gen_batch_initial_conditions(
                     acq_function=MockAcquisitionFunction(),
                     bounds=bounds,
@@ -300,16 +304,12 @@ def test_gen_batch_initial_conditions_warning(self):
                     raw_samples=10,
                     options={"seed": 1234},
                 )
-                self.assertEqual(len(ws), 1)
-                self.assertTrue(
-                    any(issubclass(w.category, BadInitialCandidatesWarning) for w in ws)
-                )
-                self.assertTrue(
-                    torch.equal(
-                        batch_initial_conditions,
-                        torch.zeros(2, 1, 2, device=self.device, dtype=dtype),
-                    )
+            self.assertTrue(
+                torch.equal(
+                    batch_initial_conditions,
+                    torch.zeros(2, 1, 2, device=self.device, dtype=dtype),
                 )
+            )
 
     def test_gen_batch_initial_conditions_transform_intra_point_constraint(self):
         for dtype in (torch.float, torch.double):
@@ -549,7 +549,10 @@ def test_gen_batch_initial_conditions_constraints(self):
                     MockAcquisitionFunction,
                     "__call__",
                     wraps=mock_acqf.__call__,
-                ) as mock_acqf_call:
+                ) as mock_acqf_call, warnings.catch_warnings():
+                    warnings.simplefilter(
+                        "ignore", category=BadInitialCandidatesWarning
+                    )
                     batch_initial_conditions = gen_batch_initial_conditions(
                         acq_function=mock_acqf,
                         bounds=bounds,
@@ -723,7 +726,10 @@ def generator(n: int, q: int, seed: int):
                     MockAcquisitionFunction,
                     "__call__",
                     wraps=mock_acqf.__call__,
-                ):
+                ), warnings.catch_warnings():
+                    warnings.simplefilter(
+                        "ignore", category=BadInitialCandidatesWarning
+                    )
                     batch_initial_conditions = gen_batch_initial_conditions(
                         acq_function=mock_acqf,
                         bounds=bounds,

test/optim/test_numpy_converter.py

@@ -230,7 +230,8 @@ def test_set_parameters(self):
 
 
 class TestScipyObjectiveAndGrad(BotorchTestCase):
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         with torch.random.fork_rng():
             torch.manual_seed(0)
             train_X = torch.linspace(0, 1, 10).unsqueeze(-1)

test/optim/utils/test_model_utils.py

@@ -116,6 +116,7 @@ def test_get_data_loader(self):
 
 class TestGetParameters(BotorchTestCase):
     def setUp(self):
+        super().setUp()
         self.module = GaussianLikelihood(
             noise_constraint=GreaterThan(1e-6, initial_value=0.123),
         )
@@ -124,7 +125,7 @@ def test_get_parameters(self):
         self.assertEqual(0, len(get_parameters(self.module, requires_grad=False)))
 
         params = get_parameters(self.module)
-        self.assertTrue(1 == len(params))
+        self.assertEqual(1, len(params))
         self.assertEqual(next(iter(params)), "noise_covar.raw_noise")
         self.assertTrue(
             self.module.noise_covar.raw_noise.equal(next(iter(params.values())))

test/test_fit.py

@@ -69,7 +69,8 @@ def __call__(self, mll, closure: Optional[Callable] = None):
 class TestFitAPI(BotorchTestCase):
     r"""Unit tests for general fitting API"""
 
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         with torch.random.fork_rng():
             torch.manual_seed(0)
             train_X = torch.linspace(0, 1, 10).unsqueeze(-1)
@@ -172,7 +173,8 @@ def mock_fit_gpytorch_mll(*args, **kwargs):
 
 
 class TestFitFallback(BotorchTestCase):
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         with torch.random.fork_rng():
             torch.manual_seed(0)
             train_X = torch.linspace(0, 1, 10).unsqueeze(-1)
@@ -377,7 +379,8 @@ def _test_exceptions(self, mll, ckpt):
 
 
 class TestFitFallbackAppoximate(BotorchTestCase):
-    def setUp(self):
+    def setUp(self) -> None:
+        super().setUp()
         with torch.random.fork_rng():
             torch.manual_seed(0)
             train_X = torch.linspace(0, 1, 10).unsqueeze(-1)

test/utils/multi_objective/box_decompositions/test_box_decomposition.py

@@ -49,6 +49,7 @@ def _get_single_cell(self):
 
 class TestBoxDecomposition(BotorchTestCase):
     def setUp(self):
+        super().setUp()
         self.ref_point_raw = torch.zeros(3, device=self.device)
         self.Y_raw = torch.tensor(
             [

test/utils/multi_objective/box_decompositions/test_utils.py

@@ -269,6 +269,7 @@ class TestFastPartitioningUtils(BotorchTestCase):
     """
 
     def setUp(self):
+        super().setUp()
         self.ref_point = -torch.tensor([10.0, 10.0, 10.0], device=self.device)
         self.U = -self.ref_point.clone().view(1, -1)
         self.Z = torch.empty(1, 3, 3, device=self.device)

test/utils/probability/test_truncated_multivariate_normal.py

@@ -28,7 +28,8 @@ def setUp(
         upper_quantile_min: float = 0.1,  # MC methods will not produce any samples.
         num_log_probs: int = 4,
         seed: int = 1,
-    ):
+    ) -> None:
+        super().setUp()
         self.seed_generator = count(seed)
         self.num_log_probs = num_log_probs
 

test/utils/test_context_managers.py

@@ -23,6 +23,7 @@
 
 class TestContextManagers(BotorchTestCase):
     def setUp(self):
+        super().setUp()
         module = self.module = Module()
         for i, name in enumerate(ascii_lowercase[:3], start=1):
             values = torch.rand(2).to(torch.float16)