add option for using posterior predictive in cross-validation

ax/modelbridge/base.py

@@ -907,12 +907,16 @@ def cross_validate(
         self,
         cv_training_data: List[Observation],
         cv_test_points: List[ObservationFeatures],
+        use_posterior_predictive: bool = False,
     ) -> List[ObservationData]:
         """Make a set of cross-validation predictions.
 
         Args:
             cv_training_data: The training data to use for cross validation.
             cv_test_points: The test points at which predictions will be made.
+            use_posterior_predictive: A boolean indicating if the predictions
+                should be from the posterior predictive (i.e. including
+                observation noise).
 
         Returns:
             A list of predictions at the test points.
@@ -936,6 +940,7 @@ def cross_validate(
                 search_space=search_space,
                 cv_training_data=cv_training_data,
                 cv_test_points=cv_test_points,
+                use_posterior_predictive=use_posterior_predictive,
             )
         # Apply reverse transforms, in reverse order
         cv_test_observations = [
@@ -952,6 +957,7 @@ def _cross_validate(
         search_space: SearchSpace,
         cv_training_data: List[Observation],
         cv_test_points: List[ObservationFeatures],
+        use_posterior_predictive: bool = False,
     ) -> List[ObservationData]:
         """Apply the terminal transform, make predictions on the test points,
         and reverse terminal transform on the results.

ax/modelbridge/cross_validation.py

@@ -78,6 +78,7 @@ def cross_validate(
     # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
     test_selector: Optional[Callable] = None,
     untransform: bool = True,
+    use_posterior_predictive: bool = False,
 ) -> List[CVResult]:
     """Cross validation for model predictions.
 
@@ -112,6 +113,12 @@ def cross_validate(
             of the original data in regions where outliers have been removed,
             we have found it to better reflect the how good the model used
             for candidate generation actually is.
+        use_posterior_predictive: A boolean indicating if the predictions
+            should be from the posterior predictive (i.e. including
+            observation noise). Note: we should reconsider how we compute
+            cross-validation and model fit metrics where there is non-
+            Gaussian noise.
+
     Returns:
         A CVResult for each observation in the training data.
     """
@@ -162,7 +169,9 @@ def cross_validate(
         # Make the prediction
         if untransform:
             cv_test_predictions = model.cross_validate(
-                cv_training_data=cv_training_data, cv_test_points=cv_test_points
+                cv_training_data=cv_training_data,
+                cv_test_points=cv_test_points,
+                use_posterior_predictive=use_posterior_predictive,
             )
         else:
             # Get test predictions in transformed space
@@ -186,6 +195,7 @@ def cross_validate(
                     search_space=search_space,
                     cv_training_data=cv_training_data,
                     cv_test_points=cv_test_points,
+                    use_posterior_predictive=use_posterior_predictive,
                 )
             # Get test observations in transformed space
             cv_test_data = deepcopy(cv_test_data)
@@ -197,7 +207,9 @@ def cross_validate(
     return result
 
 
-def cross_validate_by_trial(model: ModelBridge, trial: int = -1) -> List[CVResult]:
+def cross_validate_by_trial(
+    model: ModelBridge, trial: int = -1, use_posterior_predictive: bool = False
+) -> List[CVResult]:
     """Cross validation for model predictions on a particular trial.
 
     Uses all of the data up until the specified trial to predict each of the
@@ -206,6 +218,9 @@ def cross_validate_by_trial(model: ModelBridge, trial: int = -1) -> List[CVResul
     Args:
         model: Fitted model (ModelBridge) to cross validate.
         trial: Trial for which predictions are evaluated.
+        use_posterior_predictive: A boolean indicating if the predictions
+            should be from the posterior predictive (i.e. including
+            observation noise).
 
     Returns:
         A CVResult for each observation in the training data.
@@ -241,7 +256,9 @@ def cross_validate_by_trial(model: ModelBridge, trial: int = -1) -> List[CVResul
             cv_test_data.append(obs)
     # Make the prediction
     cv_test_predictions = model.cross_validate(
-        cv_training_data=cv_training_data, cv_test_points=cv_test_points
+        cv_training_data=cv_training_data,
+        cv_test_points=cv_test_points,
+        use_posterior_predictive=use_posterior_predictive,
     )
     # Form CVResult objects
     result = [

ax/modelbridge/discrete.py

@@ -191,6 +191,7 @@ def _cross_validate(
         search_space: SearchSpace,
         cv_training_data: List[Observation],
         cv_test_points: List[ObservationFeatures],
+        use_posterior_predictive: bool = False,
     ) -> List[ObservationData]:
         """Make predictions at cv_test_points using only the data in obs_feats
         and obs_data.
@@ -208,7 +209,11 @@ def _cross_validate(
         ]
         # Use the model to do the cross validation
         f_test, cov_test = self.model.cross_validate(
-            Xs_train=Xs_train, Ys_train=Ys_train, Yvars_train=Yvars_train, X_test=X_test
+            Xs_train=Xs_train,
+            Ys_train=Ys_train,
+            Yvars_train=Yvars_train,
+            X_test=X_test,
+            use_posterior_predictive=use_posterior_predictive,
         )
         # Convert array back to ObservationData
         return array_to_observation_data(f=f_test, cov=cov_test, outcomes=self.outcomes)

ax/modelbridge/map_torch.py

@@ -279,6 +279,7 @@ def _cross_validate(
         cv_training_data: List[Observation],
         cv_test_points: List[ObservationFeatures],
         parameters: Optional[List[str]] = None,
+        use_posterior_predictive: bool = False,
         **kwargs: Any,
     ) -> List[ObservationData]:
         """Make predictions at cv_test_points using only the data in obs_feats
@@ -294,6 +295,7 @@ def _cross_validate(
             cv_training_data=cv_training_data,
             cv_test_points=cv_test_points,
             parameters=parameters,  # we pass the map_keys too by default
+            use_posterior_predictive=use_posterior_predictive,
             **kwargs,
         )
         observation_features, observation_data = separate_observations(cv_training_data)

ax/modelbridge/random.py

@@ -100,6 +100,7 @@ def _cross_validate(
         search_space: SearchSpace,
         cv_training_data: List[Observation],
         cv_test_points: List[ObservationFeatures],
+        use_posterior_predictive: bool = False,
     ) -> List[ObservationData]:
         raise NotImplementedError
 

ax/modelbridge/tests/test_base_modelbridge.py

@@ -267,9 +267,24 @@ def warn_and_return_mock_obs(
             search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
             cv_training_data=[get_observation2trans()],
             cv_test_points=[get_observation1().features],  # untransformed after
+            use_posterior_predictive=False,
         )
         self.assertTrue(cv_predictions == [get_observation1().data])
 
+        # Test use_posterior_predictive in CV
+        modelbridge.cross_validate(
+            cv_training_data=cv_training_data,
+            cv_test_points=cv_test_points,
+            use_posterior_predictive=True,
+        )
+
+        modelbridge._cross_validate.assert_called_with(
+            search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
+            cv_training_data=[get_observation2trans()],
+            cv_test_points=[get_observation1().features],  # untransformed after
+            use_posterior_predictive=True,
+        )
+
         # Test stored training data
         obs = modelbridge.get_training_data()
         self.assertTrue(obs == [get_observation1(), get_observation2()])

ax/modelbridge/tests/test_cross_validation.py

@@ -201,7 +201,9 @@ def test_CrossValidate(self) -> None:
         # Test ModelBridge._cross_validate was called correctly.
         z = ma._cross_validate.mock_calls
         self.assertEqual(len(z), 3)
-        ma._cross_validate.assert_called_with(**self.transformed_cv_input_dict)
+        ma._cross_validate.assert_called_with(
+            **self.transformed_cv_input_dict, use_posterior_predictive=False
+        )
 
         # Test selector
 
@@ -219,6 +221,21 @@ def test_selector(obs: Observation) -> bool:
         )
         self.assertTrue(np.array_equal(sorted(all_test), np.array([2.0, 2.0, 3.0])))
 
+        # test observation noise
+        for untransform in (True, False):
+            result = cross_validate(
+                model=ma,
+                folds=-1,
+                use_posterior_predictive=True,
+                untransform=untransform,
+            )
+            if untransform:
+                mock_cv = ma.cross_validate
+            else:
+                mock_cv = ma._cross_validate
+            call_kwargs = mock_cv.mock_calls[-1].kwargs
+            self.assertTrue(call_kwargs["use_posterior_predictive"])
+
     def test_CrossValidateByTrial(self) -> None:
         # With only 1 trial
         ma = mock.MagicMock()
@@ -261,6 +278,15 @@ def test_CrossValidateByTrial(self) -> None:
         self.assertEqual(len(result), 1)
         self.assertEqual(result[0].observed.features.trial_index, 2)
 
+        mock_cv = ma.cross_validate
+        call_kwargs = mock_cv.mock_calls[-1].kwargs
+        self.assertFalse(call_kwargs["use_posterior_predictive"])
+
+        # test observation noise
+        result = cross_validate_by_trial(model=ma, use_posterior_predictive=True)
+        call_kwargs = mock_cv.mock_calls[-1].kwargs
+        self.assertTrue(call_kwargs["use_posterior_predictive"])
+
     def test_cross_validate_gives_a_useful_error_for_model_with_no_data(self) -> None:
         exp = get_branin_experiment()
         sobol = Models.SOBOL(experiment=exp, search_space=exp.search_space)

ax/modelbridge/torch.py

@@ -428,6 +428,7 @@ def _cross_validate(
         cv_training_data: List[Observation],
         cv_test_points: List[ObservationFeatures],
         parameters: Optional[List[str]] = None,
+        use_posterior_predictive: bool = False,
         **kwargs: Any,
     ) -> List[ObservationData]:
         """Make predictions at cv_test_points using only the data in obs_feats
@@ -453,6 +454,7 @@ def _cross_validate(
             datasets=datasets,
             X_test=torch.as_tensor(X_test, dtype=self.dtype, device=self.device),
             search_space_digest=search_space_digest,
+            use_posterior_predictive=use_posterior_predictive,
             **kwargs,
         )
         # Convert array back to ObservationData

ax/models/discrete_base.py

@@ -102,6 +102,7 @@ def cross_validate(
         Ys_train: List[List[float]],
         Yvars_train: List[List[float]],
         X_test: List[TParamValueList],
+        use_posterior_predictive: bool = False,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Do cross validation with the given training and test sets.
 
@@ -116,6 +117,9 @@ def cross_validate(
                 each outcome.
             Yvars_train: The variances of each entry in Ys, same shape.
             X_test: List of the j parameterizations at which to make predictions.
+            use_posterior_predictive: A boolean indicating if the predictions
+                should be from the posterior predictive (i.e. including
+                observation noise).
 
         Returns:
             2-element tuple containing

ax/models/torch/botorch.py

@@ -64,7 +64,7 @@
     ],
     Model,
 ]
-TModelPredictor = Callable[[Model, Tensor], Tuple[Tensor, Tensor]]
+TModelPredictor = Callable[[Model, Tensor, bool], Tuple[Tensor, Tensor]]
 
 
 # pyre-fixme[33]: Aliased annotation cannot contain `Any`.
@@ -466,6 +466,7 @@ def cross_validate(  # pyre-ignore [14]: `search_space_digest` arg not needed he
         self,
         datasets: List[SupervisedDataset],
         X_test: Tensor,
+        use_posterior_predictive: bool = False,
         **kwargs: Any,
     ) -> Tuple[Tensor, Tensor]:
         if self._model is None:
@@ -488,7 +489,10 @@ def cross_validate(  # pyre-ignore [14]: `search_space_digest` arg not needed he
             use_loocv_pseudo_likelihood=self.use_loocv_pseudo_likelihood,
             **self._kwargs,
         )
-        return self.model_predictor(model=model, X=X_test)  # pyre-ignore: [28]
+        # pyre-ignore: [28]
+        return self.model_predictor(
+            model=model, X=X_test, use_posterior_predictive=use_posterior_predictive
+        )
 
     def feature_importances(self) -> np.ndarray:
         return get_feature_importances_from_botorch_model(model=self._model)

ax/models/torch/botorch_modular/model.py

@@ -396,10 +396,15 @@ def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
         return f, cov
 
     def predict_from_surrogate(
-        self, surrogate_label: str, X: Tensor
+        self,
+        surrogate_label: str,
+        X: Tensor,
+        use_posterior_predictive: bool = False,
     ) -> Tuple[Tensor, Tensor]:
         """Predict from the Surrogate with the given label."""
-        return self.surrogates[surrogate_label].predict(X=X)
+        return self.surrogates[surrogate_label].predict(
+            X=X, use_posterior_predictive=use_posterior_predictive
+        )
 
     @copy_doc(TorchModel.gen)
     def gen(
@@ -504,6 +509,7 @@ def cross_validate(
         datasets: Sequence[SupervisedDataset],
         X_test: Tensor,
         search_space_digest: SearchSpaceDigest,
+        use_posterior_predictive: bool = False,
         **additional_model_inputs: Any,
     ) -> Tuple[Tensor, Tensor]:
         # Will fail if metric_names exist across multiple models
@@ -561,7 +567,9 @@ def cross_validate(
                 **additional_model_inputs,
             )
             X_test_prediction = self.predict_from_surrogate(
-                surrogate_label=surrogate_label, X=X_test
+                surrogate_label=surrogate_label,
+                X=X_test,
+                use_posterior_predictive=use_posterior_predictive,
             )
         finally:
             # Reset the surrogates back to this model's surrogate, make

ax/models/torch/botorch_modular/surrogate.py

@@ -594,17 +594,24 @@ def _discard_cached_model_and_data_if_search_space_digest_changed(
             self._last_datasets = {}
         self._last_search_space_digest = search_space_digest
 
-    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
+    def predict(
+        self, X: Tensor, use_posterior_predictive: bool = False
+    ) -> Tuple[Tensor, Tensor]:
         """Predicts outcomes given an input tensor.
 
         Args:
             X: A ``n x d`` tensor of input parameters.
+            use_posterior_predictive: A boolean indicating if the predictions
+                should be from the posterior predictive (i.e. including
+                observation noise).
 
         Returns:
             Tensor: The predicted posterior mean as an ``n x o``-dim tensor.
             Tensor: The predicted posterior covariance as a ``n x o x o``-dim tensor.
         """
-        return predict_from_model(model=self.model, X=X)
+        return predict_from_model(
+            model=self.model, X=X, use_posterior_predictive=use_posterior_predictive
+        )
 
     def best_in_sample_point(
         self,

ax/models/torch/randomforest.py

@@ -73,6 +73,7 @@ def cross_validate(  # pyre-ignore [14]: not using metric_names or ssd
         self,
         datasets: List[SupervisedDataset],
         X_test: Tensor,
+        use_posterior_predictive: bool = False,
     ) -> Tuple[Tensor, Tensor]:
         Xs, Ys, Yvars = _datasets_to_legacy_inputs(datasets=datasets)
         cv_models: List[RandomForestRegressor] = []