Suppress warnings in PairwiseGP. (#428)

Summary:
Suppresses `RuntimeWarning` within the `fsolve` calls.
Also, the warning in `_add_jitter` was not emitting a warning of type `RuntimeWarning`, so it was not properly suppressed in the tests.

Pull Request resolved: #428

Test Plan: unit tests

Reviewed By: danielrjiang

Differential Revision: D21247256

Pulled By: Balandat

fbshipit-source-id: 9deb70317cdfbd27c895cf46135dae70be720c90

Author: Balandat

botorch/models/pairwise_gp.py

@@ -149,7 +149,7 @@ def __init__(
 
         self.to(self.datapoints)
 
-    def __deepcopy__(self, memo):
+    def __deepcopy__(self, memo) -> PairwiseGP:
         attrs = (
             "datapoints",
             "comparisons",
@@ -260,27 +260,28 @@ def _prior_predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
         pred_covar = self._calc_covar(X, X)
         return pred_mean, pred_covar
 
-    def _add_jitter(self, X) -> Tensor:
-        X = X.clone()
+    def _add_jitter(self, X: Tensor) -> Tensor:
         jitter_prev = 0
+        Eye = torch.eye(X.size(-1)).expand(X.shape)
         for i in range(3):
             jitter_new = self._jitter * (10 ** i)
-            X = X + torch.eye(X.size(-1)).expand(X.shape) * (jitter_new - jitter_prev)
+            X = X + (jitter_new - jitter_prev) * Eye
             jitter_prev = jitter_new
             # This may be VERY slow given upstream pytorch issue:
             # https://github.com/pytorch/pytorch/issues/34272
             try:
                 _ = torch.cholesky(X)
                 warnings.warn(
-                    f"X is not a p.d. matrix; "
-                    f"Added jitter of {jitter_new} to the diagonal",
+                    "X is not a p.d. matrix; "
+                    f"Added jitter of {jitter_new:.2e} to the diagonal",
                     RuntimeWarning,
                 )
                 return X
             except RuntimeError:
                 continue
         warnings.warn(
-            f"Failed to turn X into p.d. after trying to add adding {jitter_new} jitter"
+            f"Failed to render X p.d. after adding {jitter_new:.2e} jitter",
+            RuntimeWarning,
         )
         return X
 
@@ -524,15 +525,17 @@ def _update(self, **kwargs) -> None:
                         ci_v[i],
                         True,
                     )
-                    x[i] = optimize.fsolve(
-                        x0=x0[i],
-                        func=self._grad_posterior_f,
-                        fprime=self._hess_posterior_f,
-                        xtol=xtol,
-                        maxfev=maxfev,
-                        args=fsolve_args,
-                        **kwargs,
-                    )
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings("ignore", category=RuntimeWarning)
+                        x[i] = optimize.fsolve(
+                            x0=x0[i],
+                            func=self._grad_posterior_f,
+                            fprime=self._hess_posterior_f,
+                            xtol=xtol,
+                            maxfev=maxfev,
+                            args=fsolve_args,
+                            **kwargs,
+                        )
                 x = x.reshape(*init_x0_size)
             else:
                 fsolve_args = (
@@ -544,15 +547,17 @@ def _update(self, **kwargs) -> None:
                     self.covar_inv,
                     True,
                 )
-                x = optimize.fsolve(
-                    x0=x0,
-                    func=self._grad_posterior_f,
-                    fprime=self._hess_posterior_f,
-                    xtol=xtol,
-                    maxfev=maxfev,
-                    args=fsolve_args,
-                    **kwargs,
-                )
+                with warnings.catch_warnings():
+                    warnings.filterwarnings("ignore", category=RuntimeWarning)
+                    x = optimize.fsolve(
+                        x0=x0,
+                        func=self._grad_posterior_f,
+                        fprime=self._hess_posterior_f,
+                        xtol=xtol,
+                        maxfev=maxfev,
+                        args=fsolve_args,
+                        **kwargs,
+                    )
 
             self._x0 = x.copy()  # save for warm-starting
             f = torch.tensor(x, **self.tkwargs)

test/models/test_pairwise_gp.py

@@ -247,7 +247,7 @@ def test_condition_on_observations(self):
                         "datapoints": model_kwargs["datapoints"][0],
                         "comparisons": model_kwargs["comparisons"][0],
                     }
-                    model_non_batch = type(model)(**model_kwargs_non_batch)
+                    model_non_batch = model.__class__(**model_kwargs_non_batch)
                     model_non_batch.load_state_dict(state_dict_non_batch)
                     model_non_batch.eval()
                     model_non_batch.posterior(