EFF Make GaussianProcessRegressor.predict faster when return_std and return_cov are false (scikit-learn#31431)

Co-authored-by: Loïc Estève <loic.esteve@ymail.com>
Co-authored-by: Jérémie du Boisberranger <jeremie@probabl.ai>

Author: 3 people

doc/whats_new/upcoming_changes/sklearn.gaussian_process/31431.efficiency.rst

@@ -0,0 +1,3 @@
+- make :class:`GaussianProcessRegressor.predict` faster when `return_cov` and
+  `return_std` are both `False`.
+  By :user:`Rafael Ayllón Gavilán <RafaAyGar>`.

sklearn/gaussian_process/_gpr.py

@@ -450,6 +450,9 @@ def predict(self, X, return_std=False, return_cov=False):
             if y_mean.ndim > 1 and y_mean.shape[1] == 1:
                 y_mean = np.squeeze(y_mean, axis=1)
 
+            if not return_cov and not return_std:
+                return y_mean
+
             # Alg 2.1, page 19, line 5 -> v = L \ K(X_test, X_train)^T
             V = solve_triangular(
                 self.L_, K_trans.T, lower=GPR_CHOLESKY_LOWER, check_finite=False
@@ -467,7 +470,7 @@ def predict(self, X, return_std=False, return_cov=False):
                     y_cov = np.squeeze(y_cov, axis=2)
 
                 return y_mean, y_cov
-            elif return_std:
+            else:  # return_std
                 # Compute variance of predictive distribution
                 # Use einsum to avoid explicitly forming the large matrix
                 # V^T @ V just to extract its diagonal afterward.
@@ -492,8 +495,6 @@ def predict(self, X, return_std=False, return_cov=False):
                     y_var = np.squeeze(y_var, axis=1)
 
                 return y_mean, np.sqrt(y_var)
-            else:
-                return y_mean
 
     def sample_y(self, X, n_samples=1, random_state=0):
         """Draw samples from Gaussian process and evaluate at X.

sklearn/gaussian_process/tests/test_gpr.py

@@ -847,3 +847,15 @@ def test_gpr_predict_input_not_modified():
     _, _ = gpr.predict(X2, return_std=True)
 
     assert_allclose(X2, X2_copy)
+
+
+@pytest.mark.parametrize("kernel", kernels)
+def test_gpr_predict_no_cov_no_std_return(kernel):
+    """
+    Check that only y_mean is returned when return_cov=False and
+    return_std=False.
+    """
+    gpr = GaussianProcessRegressor(kernel=kernel).fit(X, y)
+    y_pred = gpr.predict(X, return_cov=False, return_std=False)
+
+    assert_allclose(y_pred, y)