Silence input normalization / standardization warnings with custom covariance modules (#2745)

Summary:
Pull Request resolved: #2745

The default covar modules are designed to work with normalized & standardized data. If the user has a custom covar module, they may need a different scaling of the data, which makes these warnings irrelevant.

This diff adds an option to restrict input data checks to NaNs only, and enables it whenever a custom covar module is given.

Reviewed By: sdaulton

Differential Revision: D69613875

fbshipit-source-id: 8545752ae4f532a76c59983b963fbd6b873dde1c

Author: saitcakmak

botorch/models/approximate_gp.py

@@ -399,7 +399,11 @@ def __init__(
                 )
                 train_Y, _ = outcome_transform(train_Y, X=transformed_X)
             self._validate_tensor_args(X=transformed_X, Y=train_Y)
-            validate_input_scaling(train_X=transformed_X, train_Y=train_Y)
+            validate_input_scaling(
+                train_X=transformed_X,
+                train_Y=train_Y,
+                check_nans_only=covar_module is not None,
+            )
             if train_Y.shape[-1] != num_outputs:
                 num_outputs = train_Y.shape[-1]
 

botorch/models/gp_regression.py

@@ -171,6 +171,7 @@ def __init__(
             train_Y=train_Y,
             train_Yvar=train_Yvar,
             ignore_X_dims=ignore_X_dims,
+            check_nans_only=covar_module is not None,
         )
         self._set_dimensions(train_X=train_X, train_Y=train_Y)
         train_X, train_Y, train_Yvar = self._transform_tensor_args(

botorch/models/utils/assorted.py

@@ -228,6 +228,7 @@ def validate_input_scaling(
     train_Yvar: Tensor | None = None,
     raise_on_fail: bool = False,
     ignore_X_dims: list[int] | None = None,
+    check_nans_only: bool = False,
 ) -> None:
     r"""Helper function to validate input data to models.
 
@@ -243,6 +244,10 @@ def validate_input_scaling(
             raised if NaN values are present).
         ignore_X_dims: For this subset of dimensions from `{1, ..., d}`, ignore the
             min-max scaling check.
+        check_nans_only: If True, only check for NaN values. Skips min-max scaling
+            and standardization checks. This is used when the model is provided
+            with a custom covariance module, to avoid potentially irrelevant
+            warnings.
 
     This function is typically called inside the constructor of standard BoTorch
     models. It validates the following:
@@ -261,10 +266,11 @@ def validate_input_scaling(
         check_no_nans(train_Yvar)
         if torch.any(train_Yvar < 0):
             raise InputDataError("Input data contains negative variances.")
-    check_min_max_scaling(
-        X=train_X, raise_on_fail=raise_on_fail, ignore_dims=ignore_X_dims
-    )
-    check_standardization(Y=train_Y, raise_on_fail=raise_on_fail)
+    if not check_nans_only:
+        check_min_max_scaling(
+            X=train_X, raise_on_fail=raise_on_fail, ignore_dims=ignore_X_dims
+        )
+        check_standardization(Y=train_Y, raise_on_fail=raise_on_fail)
 
 
 def mod_batch_shape(module: Module, names: list[str], b: int) -> None:

test/models/utils/test_assorted.py

@@ -87,7 +87,7 @@ def test_check_min_max_scaling(self):
         X = 0.1 + 0.8 * torch.rand(4, 2, 3)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X)
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
         check_min_max_scaling(X=X, raise_on_fail=True)
         with self.assertWarnsRegex(
             expected_warning=InputDataWarning, expected_regex="not scaled"
@@ -100,30 +100,34 @@ def test_check_min_max_scaling(self):
         Xstd = (X - Xmin) / (Xmax - Xmin)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=Xstd)
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
         check_min_max_scaling(X=Xstd, raise_on_fail=True)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=Xstd, strict=True)
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
         check_min_max_scaling(X=Xstd, strict=True, raise_on_fail=True)
         # check violation
         X[0, 0, 0] = 2
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X)
-            self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
-            self.assertTrue(any("not contained" in str(w.message) for w in ws))
+        self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertTrue(any("not contained" in str(w.message) for w in ws))
         with self.assertRaises(InputDataError):
             check_min_max_scaling(X=X, raise_on_fail=True)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X, strict=True)
-            self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
-            self.assertTrue(any("not contained" in str(w.message) for w in ws))
+        self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertTrue(any("not contained" 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 ignore_dims
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X, ignore_dims=[0])
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        # all dims ignored
+        with warnings.catch_warnings(record=True) as ws:
+            check_min_max_scaling(X=X, ignore_dims=[0, 1, 2])
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
 
     def test_check_standardization(self):
         # Ensure that it is not filtered out.
@@ -181,6 +185,11 @@ def test_validate_input_scaling(self):
         # check that errors are raised when requested
         with self.assertRaises(InputDataError):
             validate_input_scaling(train_X=train_X, train_Y=train_Y, raise_on_fail=True)
+        # check that normalization & standardization checks & errors are skipped when
+        # check_nans_only is True
+        validate_input_scaling(
+            train_X=train_X, train_Y=train_Y, raise_on_fail=True, check_nans_only=True
+        )
         # check that no errors are being raised if everything is standardized
         train_X_min = train_X.min(dim=-1, keepdim=True)[0]
         train_X_max = train_X.max(dim=-1, keepdim=True)[0]
@@ -202,6 +211,11 @@ def test_validate_input_scaling(self):
         train_X_std[0, 0, 0] = float("nan")
         with self.assertRaises(InputDataError):
             validate_input_scaling(train_X=train_X_std, train_Y=train_Y_std)
+        # NaNs still raise errors when check_nans_only is True
+        with self.assertRaises(InputDataError):
+            validate_input_scaling(
+                train_X=train_X_std, train_Y=train_Y_std, check_nans_only=True
+            )
 
 
 class TestGPTPosteriorSettings(BotorchTestCase):