Improve Documentation of Testing Base Classes in BoTorch (#2770)

Summary:
Pull Request resolved: #2770

I added doc strings (and some abstract methods).

I came across the following issues doing so:

- There are multiple private functions, like `_get_random_data`. These are used outside, though, thus the editor shows them as unused. Should they maybe not be private instead?
- The double naming of `BaseTestProblem` and `functions` is confusing for me. Is there a reason for this? Should we maybe change it to only use one of the two?
- Only `TestCorruptedProblemsMixin.rosenbrock_problem` seems to be used from that Mixin. Should we maybe change this to not be a Mixin, but just the rosenbrock variable instad, or not have the `outlier_generator` inside the Mixin?
- I think we should delete `rsample_from_base_samples`. It is not used (if my command+click works here properly) and hard to understand.

Reviewed By: esantorella

Differential Revision: D71198064

fbshipit-source-id: 9dd660045aaa3424f09063be0efb747b5cfc321d

Author: SamuelGabriel

botorch/utils/testing.py

@@ -28,7 +28,7 @@
 from botorch.models.transforms.input import Normalize
 from botorch.posteriors import GPyTorchPosterior
 from botorch.utils import apply_constraints
-from botorch.utils.testing import _get_test_posterior, BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, get_test_posterior
 from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
 from linear_operator.operators.dense_linear_operator import to_linear_operator
 
@@ -67,7 +67,7 @@ def test_scalarized_posterior_transform(self):
             offset = torch.rand(1).item()
             weights = torch.randn(m, device=self.device, dtype=dtype)
             obj = ScalarizedPosteriorTransform(weights=weights, offset=offset)
-            posterior = _get_test_posterior(
+            posterior = get_test_posterior(
                 batch_shape, m=m, device=self.device, dtype=dtype
             )
             mean, covar = (

test/acquisition/test_objective.py

@@ -19,7 +19,7 @@
 from botorch.sampling import SobolQMCNormalSampler
 from botorch.utils.datasets import SupervisedDataset
 from botorch.utils.test_helpers import get_pvar_expected
-from botorch.utils.testing import _get_random_data, BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, get_random_data
 from gpytorch.kernels import RBFKernel
 from gpytorch.likelihoods import FixedNoiseGaussianLikelihood, GaussianLikelihood
 from gpytorch.means import ConstantMean, ZeroMean
@@ -38,7 +38,7 @@ def _get_model_and_data(
         **tkwargs,
     ):
         extra_model_kwargs = extra_model_kwargs or {}
-        train_X, train_Y = _get_random_data(batch_shape=batch_shape, m=m, **tkwargs)
+        train_X, train_Y = get_random_data(batch_shape=batch_shape, m=m, **tkwargs)
         model_kwargs = {
             "train_X": train_X,
             "train_Y": train_Y,
@@ -182,7 +182,7 @@ def test_default_transforms(self):
             ("Default", "None", "Log"),  # Outcome transform
         ):
             tkwargs = {"device": self.device, "dtype": dtype}
-            train_X, train_Y = _get_random_data(batch_shape=batch_shape, m=m, **tkwargs)
+            train_X, train_Y = get_random_data(batch_shape=batch_shape, m=m, **tkwargs)
 
             model_kwargs = {}
             if octf == "None":
@@ -241,7 +241,7 @@ def test_condition_on_observations(self):
             # test condition_on_observations
             fant_shape = torch.Size([2])
             # fantasize at different input points
-            X_fant, Y_fant = _get_random_data(
+            X_fant, Y_fant = get_random_data(
                 batch_shape=fant_shape + batch_shape, m=m, n=3, **tkwargs
             )
             c_kwargs = (
@@ -459,7 +459,7 @@ def _get_model_and_data(
         **tkwargs,
     ):
         extra_model_kwargs = extra_model_kwargs or {}
-        train_X, train_Y = _get_random_data(batch_shape=batch_shape, m=m, **tkwargs)
+        train_X, train_Y = get_random_data(batch_shape=batch_shape, m=m, **tkwargs)
         model_kwargs = {
             "train_X": train_X,
             "train_Y": train_Y,

test/models/test_gp_regression.py

@@ -16,7 +16,7 @@
 from botorch.posteriors import GPyTorchPosterior
 from botorch.sampling import SobolQMCNormalSampler
 from botorch.utils.datasets import SupervisedDataset
-from botorch.utils.testing import _get_random_data, BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, get_random_data
 from gpytorch.kernels.scale_kernel import ScaleKernel
 from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
 from gpytorch.means import ConstantMean
@@ -30,7 +30,7 @@ def _get_random_data_with_fidelity(
     r"""Construct test data.
     For this test, by convention the trailing dimensions are the fidelity dimensions
     """
-    train_x, train_y = _get_random_data(
+    train_x, train_y = get_random_data(
         batch_shape=batch_shape, m=m, d=d, n=n, **tkwargs
     )
     s = torch.rand(n, n_fidelity, **tkwargs).repeat(batch_shape + torch.Size([1, 1]))

test/models/test_gp_regression_fidelity.py

@@ -18,7 +18,7 @@
 from botorch.sampling import SobolQMCNormalSampler
 from botorch.utils.datasets import SupervisedDataset
 from botorch.utils.test_helpers import get_pvar_expected
-from botorch.utils.testing import _get_random_data, BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, get_random_data
 from gpytorch.kernels.kernel import AdditiveKernel, ProductKernel
 from gpytorch.kernels.rbf_kernel import RBFKernel
 from gpytorch.kernels.scale_kernel import ScaleKernel
@@ -44,7 +44,7 @@ def test_gp(self):
             # to test without that transform we need to explicitly pass in `None`.
             outcome_transform_kwargs = {} if use_octf else {"outcome_transform": None}
 
-            train_X, train_Y = _get_random_data(
+            train_X, train_Y = get_random_data(
                 batch_shape=batch_shape, m=m, d=d, **tkwargs
             )
             cat_dims = list(range(ncat))
@@ -151,7 +151,7 @@ def test_condition_on_observations__(self):
             (torch.Size([2]), 1, 2, torch.double, False),
         ):
             tkwargs = {"device": self.device, "dtype": dtype}
-            train_X, train_Y = _get_random_data(
+            train_X, train_Y = get_random_data(
                 batch_shape=batch_shape, m=m, d=d, **tkwargs
             )
             cat_dims = list(range(ncat))
@@ -169,7 +169,7 @@ def test_condition_on_observations__(self):
             fant_shape = torch.Size([2])
 
             # fantasize at different input points
-            X_fant, Y_fant = _get_random_data(
+            X_fant, Y_fant = get_random_data(
                 fant_shape + batch_shape, m=m, d=d, n=3, **tkwargs
             )
             additional_kwargs = (
@@ -255,7 +255,7 @@ def test_fantasize(self):
             (torch.Size([2]), 1, 2, torch.double, False),
         ):
             tkwargs = {"device": self.device, "dtype": dtype}
-            train_X, train_Y = _get_random_data(
+            train_X, train_Y = get_random_data(
                 batch_shape=batch_shape, m=m, d=d, **tkwargs
             )
             train_Yvar = torch.full_like(train_Y, 0.1) if observed_noise else None
@@ -281,7 +281,7 @@ def test_subset_model(self):
             (torch.float, torch.double),
         ):
             tkwargs = {"device": self.device, "dtype": dtype}
-            train_X, train_Y = _get_random_data(
+            train_X, train_Y = get_random_data(
                 batch_shape=batch_shape, m=m, d=d, **tkwargs
             )
             cat_dims = list(range(ncat))
@@ -307,7 +307,7 @@ def test_construct_inputs(self):
             (torch.Size(), torch.Size([2])), (1, 2), (torch.float, torch.double)
         ):
             tkwargs = {"device": self.device, "dtype": dtype}
-            X, Y = _get_random_data(batch_shape=batch_shape, m=1, d=d, **tkwargs)
+            X, Y = get_random_data(batch_shape=batch_shape, m=1, d=d, **tkwargs)
             cat_dims = list(range(ncat))
             training_data = SupervisedDataset(
                 X,

test/models/test_gp_regression_mixed.py

@@ -29,7 +29,7 @@
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.sampling.normal import SobolQMCNormalSampler
 from botorch.utils.test_helpers import SimpleGPyTorchModel
-from botorch.utils.testing import _get_random_data, BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, get_random_data
 from gpytorch import ExactMarginalLogLikelihood
 from gpytorch.distributions import MultivariateNormal
 from gpytorch.kernels import RBFKernel, ScaleKernel
@@ -492,7 +492,7 @@ def test_model_list_gpytorch_model(self):
             self.assertIsInstance(posterior, GPyTorchPosterior)
             self.assertEqual(posterior.mean.shape, torch.Size([2, 2]))
             # test multioutput
-            train_x_raw, train_y = _get_random_data(
+            train_x_raw, train_y = get_random_data(
                 batch_shape=torch.Size(), m=1, n=10, **tkwargs
             )
             task_idx = torch.cat(

test/models/test_gpytorch.py

@@ -14,7 +14,7 @@
 from botorch.fit import fit_gpytorch_mll
 from botorch.models.latent_kronecker_gp import LatentKroneckerGP, MinMaxStandardize
 from botorch.models.transforms import Normalize
-from botorch.utils.testing import _get_random_data, BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, get_random_data
 from botorch.utils.types import DEFAULT
 from gpytorch.kernels import MaternKernel, RBFKernel, ScaleKernel
 from gpytorch.likelihoods import FixedNoiseGaussianLikelihood, GaussianLikelihood
@@ -27,7 +27,7 @@
 def _get_data_with_missing_entries(
     n_train: int, d: int, m: int, batch_shape: torch.Size, tkwargs: dict
 ):
-    train_X, train_Y = _get_random_data(
+    train_X, train_Y = get_random_data(
         batch_shape=batch_shape, m=m, d=d, n=n_train, **tkwargs
     )
 
@@ -586,7 +586,7 @@ def test_iterative_methods(self):
             (torch.float, torch.double),  # dtype
         ):
             tkwargs = {"device": self.device, "dtype": dtype}
-            train_X, train_Y = _get_random_data(
+            train_X, train_Y = get_random_data(
                 batch_shape=batch_shape, m=m, d=d, n=n_train, **tkwargs
             )
 
@@ -619,7 +619,7 @@ def test_iterative_methods(self):
     def test_not_implemented(self):
         batch_shape = torch.Size([])
         tkwargs = {"device": self.device, "dtype": torch.double}
-        train_X, train_Y = _get_random_data(batch_shape=batch_shape, m=1, **tkwargs)
+        train_X, train_Y = get_random_data(batch_shape=batch_shape, m=1, **tkwargs)
 
         model = LatentKroneckerGP(
             train_X=train_X,

test/models/test_latent_kronecker_gp.py

@@ -23,7 +23,7 @@
 from botorch.sampling.base import MCSampler
 from botorch.sampling.list_sampler import ListSampler
 from botorch.sampling.normal import IIDNormalSampler
-from botorch.utils.testing import _get_random_data, BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, get_random_data
 from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
 from gpytorch.kernels import RBFKernel
 from gpytorch.likelihoods import LikelihoodList
@@ -41,13 +41,9 @@
 def _get_model(
     fixed_noise=False, outcome_transform: str = "None", use_intf=False, **tkwargs
 ) -> ModelListGP:
-    train_x1, train_y1 = _get_random_data(
-        batch_shape=torch.Size(), m=1, n=10, **tkwargs
-    )
+    train_x1, train_y1 = get_random_data(batch_shape=torch.Size(), m=1, n=10, **tkwargs)
     train_y1 = torch.exp(train_y1)
-    train_x2, train_y2 = _get_random_data(
-        batch_shape=torch.Size(), m=1, n=11, **tkwargs
-    )
+    train_x2, train_y2 = get_random_data(batch_shape=torch.Size(), m=1, n=11, **tkwargs)
     if outcome_transform == "Standardize":
         octfs = [Standardize(m=1), Standardize(m=1)]
     elif outcome_transform == "Log":
@@ -280,7 +276,7 @@ def test_ModelListGP_fixed_noise(self) -> None:
 
     def test_ModelListGP_single(self):
         tkwargs = {"device": self.device, "dtype": torch.float}
-        train_x1, train_y1 = _get_random_data(
+        train_x1, train_y1 = get_random_data(
             batch_shape=torch.Size(), m=1, n=10, **tkwargs
         )
         model1 = SingleTaskGP(train_X=train_x1, train_Y=train_y1)
@@ -296,7 +292,7 @@ def test_ModelListGP_multi_task(self, use_outcome_transform: bool = False):
         outcome_transform_kwargs = (
             {} if use_outcome_transform else {"outcome_transform": None}
         )
-        train_x_raw, train_y = _get_random_data(
+        train_x_raw, train_y = get_random_data(
             batch_shape=torch.Size(), m=1, n=10, **tkwargs
         )
         task_idx = torch.cat(

test/models/test_model_list_gp_regression.py

@@ -48,9 +48,9 @@
 from botorch.sampling.normal import IIDNormalSampler
 from botorch.utils.sampling import manual_seed, unnormalize
 from botorch.utils.testing import (
-    _get_max_violation_of_bounds,
-    _get_max_violation_of_constraints,
     BotorchTestCase,
+    get_max_violation_of_bounds,
+    get_max_violation_of_constraints,
     MockAcquisitionFunction,
     MockModel,
     MockPosterior,
@@ -61,11 +61,11 @@ class TestBoundsAndConstraintCheckers(BotorchTestCase):
     def test_bounds_check(self) -> None:
         bounds = torch.tensor([[1, 2], [3, 4]], device=self.device)
         samples = torch.tensor([[2, 3], [2, 3.1]], device=self.device)[None, :, :]
-        result = _get_max_violation_of_bounds(samples, bounds)
+        result = get_max_violation_of_bounds(samples, bounds)
         self.assertAlmostEqual(result, -0.9, delta=1e-6)
 
         samples = torch.tensor([[2, 3], [2, 4.1]], device=self.device)[None, :, :]
-        result = _get_max_violation_of_bounds(samples, bounds)
+        result = get_max_violation_of_bounds(samples, bounds)
         self.assertAlmostEqual(result, 0.1, delta=1e-6)
 
     def test_constraint_check(self) -> None:
@@ -77,10 +77,10 @@ def test_constraint_check(self) -> None:
             )
         ]
         samples = torch.tensor([[2, 3], [2, 3.1]], device=self.device)[None, :, :]
-        result = _get_max_violation_of_constraints(samples, constraints, equality=True)
+        result = get_max_violation_of_constraints(samples, constraints, equality=True)
         self.assertAlmostEqual(result, 0.1, delta=1e-6)
 
-        result = _get_max_violation_of_constraints(samples, constraints, equality=False)
+        result = get_max_violation_of_constraints(samples, constraints, equality=False)
         self.assertAlmostEqual(result, 0.0, delta=1e-6)
 
 
@@ -268,7 +268,7 @@ def test_gen_batch_initial_conditions(self):
             self.assertEqual(batch_initial_conditions.device, bounds.device)
             self.assertEqual(batch_initial_conditions.dtype, bounds.dtype)
             self.assertLess(
-                _get_max_violation_of_bounds(batch_initial_conditions, bounds),
+                get_max_violation_of_bounds(batch_initial_conditions, bounds),
                 1e-6,
             )
             batch_shape = (
@@ -347,7 +347,7 @@ def test_gen_batch_initial_conditions_topn(self):
             self.assertEqual(batch_initial_conditions.device, bounds.device)
             self.assertEqual(batch_initial_conditions.dtype, bounds.dtype)
             self.assertLess(
-                _get_max_violation_of_bounds(batch_initial_conditions, bounds),
+                get_max_violation_of_bounds(batch_initial_conditions, bounds),
                 1e-6,
             )
             batch_shape = (
@@ -409,7 +409,7 @@ def test_gen_batch_initial_conditions_highdim(self):
             self.assertEqual(batch_initial_conditions.device, bounds.device)
             self.assertEqual(batch_initial_conditions.dtype, bounds.dtype)
             self.assertLess(
-                _get_max_violation_of_bounds(batch_initial_conditions, bounds), 1e-6
+                get_max_violation_of_bounds(batch_initial_conditions, bounds), 1e-6
             )
             if ffs is not None:
                 for idx, val in ffs.items():
@@ -637,18 +637,18 @@ def _to_self_device(
                     return None if x is None else x.to(device=self.device)
 
                 self.assertLess(
-                    _get_max_violation_of_bounds(_to_self_device(samples), bounds), tol
+                    get_max_violation_of_bounds(_to_self_device(samples), bounds), tol
                 )
 
                 self.assertLess(
-                    _get_max_violation_of_constraints(
+                    get_max_violation_of_constraints(
                         _to_self_device(samples), constraints=equalities, equality=True
                     ),
                     tol,
                 )
 
                 self.assertLess(
-                    _get_max_violation_of_constraints(
+                    get_max_violation_of_constraints(
                         _to_self_device(samples),
                         constraints=inequalities,
                         equality=False,
@@ -708,19 +708,19 @@ def test_gen_batch_initial_conditions_constraints(self):
                 self.assertEqual(batch_initial_conditions.device, bounds.device)
                 self.assertEqual(batch_initial_conditions.dtype, bounds.dtype)
                 self.assertLess(
-                    _get_max_violation_of_bounds(batch_initial_conditions, bounds),
+                    get_max_violation_of_bounds(batch_initial_conditions, bounds),
                     1e-6,
                 )
                 self.assertLess(
-                    _get_max_violation_of_constraints(
+                    get_max_violation_of_constraints(
                         batch_initial_conditions,
                         inequality_constraints,
                         equality=False,
                     ),
                     1e-6,
                 )
                 self.assertLess(
-                    _get_max_violation_of_constraints(
+                    get_max_violation_of_constraints(
                         batch_initial_conditions,
                         equality_constraints,
                         equality=True,
@@ -821,7 +821,7 @@ def test_gen_batch_initial_conditions_interpoint_constraints(self):
                     batch_initial_conditions[1, 2, 0],
                 )
                 self.assertLess(
-                    _get_max_violation_of_constraints(
+                    get_max_violation_of_constraints(
                         batch_initial_conditions,
                         inequality_constraints,
                         equality=False,
@@ -886,7 +886,7 @@ def generator(n: int, q: int, seed: int | None):
                 self.assertEqual(batch_initial_conditions.dtype, bounds.dtype)
                 self.assertTrue((batch_initial_conditions[..., -1] == 0.42).all())
                 self.assertLess(
-                    _get_max_violation_of_bounds(batch_initial_conditions, bounds),
+                    get_max_violation_of_bounds(batch_initial_conditions, bounds),
                     1e-6,
                 )
                 if ffs is not None:
@@ -981,7 +981,7 @@ def test_gen_batch_initial_conditions_fixed_X_fantasies(self):
                 self.assertEqual(batch_initial_conditions.device, bounds.device)
                 self.assertEqual(batch_initial_conditions.dtype, bounds.dtype)
                 self.assertLess(
-                    _get_max_violation_of_bounds(batch_initial_conditions, bounds),
+                    get_max_violation_of_bounds(batch_initial_conditions, bounds),
                     1e-6,
                 )
                 batch_shape = (