Repository: scikit-learn/scikit-learn
Pull request: FIX add `n_targets` allowing consistent shape prediction before calling `fit` in GPR
Fixes #22430

<!--
Thanks for contributing a pull request! Please ensure you have taken a look at
the contribution guidelines: https://github.com/scikit-learn/scikit-learn/blob/main/CONTRIBUTING.md
-->

#### Reference Issues/PRs
<!--
Example: Fixes #1234. See also #3456.
Please use keywords (e.g., Fixes) to create link to the issues or pull requests
you resolved, so that they will automatically be closed when your pull request
is merged. See https://github.com/blog/1506-closing-issues-via-pull-requests
-->
This PR adds a new argument `n_targets` for `GaussianProcessRegressor` as suggested in issue #22430 , which will be used to determine the output shape of both `predict` and `sample_y`, so the shape will be consistent before and after calling `fit` method.

#### Any other comments?
I'm not sure if we should raise error during the fit process if the dimension of`y`  does not match `n_targets`, since once we fit the model the output shape will be determined by the shape of y.

<!--
Please be aware that we are a loose team of volunteers so patience is
necessary; assistance handling other issues is very welcome. We value
all user contributions, no matter how minor they are. If we are slow to
review, either the pull request needs some benchmarking, tinkering,
convincing, etc. or more likely the reviewers are simply busy. In either
case, we ask for your understanding during the review process.
For more information, see our FAQ on this topic:
http://scikit-learn.org/dev/faq.html#why-is-my-pull-request-not-getting-any-attention.

Thanks for contributing!
-->


Problem statement:
GPR `sample_y` enforce `n_targets=1` before calling `fit`
In `GaussianProcessRegressor`, sampling in the prior before calling `fit` via `sample_y` will assume that `y` is made of a single target. However, this is not necessarily the case. Therefore, the shape of the output of `sample_y` before and after `fit` is different.

In order to solve this inconsistency, we need to introduce a new parameter `n_targets=None`. Before calling `fit` this parameter should be explicitly set by the user. After `fit`, we can use the information of the target seen during `fit` without explicitly setting the parameter.

Hints from issue discussion:
I see that we have the same issue with `predict` indeed.

Commit message: fix test case

Patch to review:
diff --git a/sklearn/gaussian_process/_gpr.py b/sklearn/gaussian_process/_gpr.py
index c7d8db7b63702..4ba094bd31fcd 100644
--- a/sklearn/gaussian_process/_gpr.py
+++ b/sklearn/gaussian_process/_gpr.py
@@ -108,6 +108,12 @@ def optimizer(obj_func, initial_theta, bounds):
         which might cause predictions to change if the data is modified
         externally.
 
+    n_targets : int, default=None
+        Dimension of the target values. Used to decide the output dimension
+        when sampling with prior distribution.
+
+        .. versionadded:: 1.1
+
     random_state : int, RandomState instance or None, default=None
         Determines random number generation used to initialize the centers.
         Pass an int for reproducible results across multiple function calls.
@@ -182,6 +188,7 @@ def __init__(
         n_restarts_optimizer=0,
         normalize_y=False,
         copy_X_train=True,
+        n_targets=None,
         random_state=None,
     ):
         self.kernel = kernel
@@ -190,6 +197,7 @@ def __init__(
         self.n_restarts_optimizer = n_restarts_optimizer
         self.normalize_y = normalize_y
         self.copy_X_train = copy_X_train
+        self.n_targets = n_targets
         self.random_state = random_state
 
     def fit(self, X, y):
@@ -378,12 +386,26 @@ def predict(self, X, return_std=False, return_cov=False):
                 )
             else:
                 kernel = self.kernel
-            y_mean = np.zeros(X.shape[0])
+
+            n_targets = self.n_targets or 1
+            if n_targets == 1:
+                y_mean = np.zeros(X.shape[0])
+            else:
+                y_mean = np.zeros(shape=(X.shape[0], self.n_targets))
+
             if return_cov:
                 y_cov = kernel(X)
+                if n_targets > 1:
+                    y_cov = np.repeat(
+                        np.expand_dims(y_cov, -1), repeats=n_targets, axis=-1
+                    )
                 return y_mean, y_cov
             elif return_std:
                 y_var = kernel.diag(X)
+                if n_targets > 1:
+                    y_var = np.repeat(
+                        np.expand_dims(y_var, -1), repeats=n_targets, axis=-1
+                    )
                 return y_mean, np.sqrt(y_var)
             else:
                 return y_mean
diff --git a/sklearn/gaussian_process/tests/test_gpr.py b/sklearn/gaussian_process/tests/test_gpr.py
index a8e6eda3bf667..9ace73ca05e82 100644
--- a/sklearn/gaussian_process/tests/test_gpr.py
+++ b/sklearn/gaussian_process/tests/test_gpr.py
@@ -764,3 +764,40 @@ def test_sample_y_shapes(normalize_y, n_targets):
 
     y_samples = model.sample_y(X_test, n_samples=n_samples_y_test)
     assert y_samples.shape == y_test_shape
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_sample_y_shape_with_prior(n_targets):
+    """Check the output shape of `sample_y` is consistent before and after `fit`"""
+    rng = np.random.RandomState(1024)
+
+    X = rng.randn(10, 3)
+    y = rng.randn(10, n_targets or 1)
+
+    model = GaussianProcessRegressor(n_targets=n_targets)
+    shape_before_fit = model.sample_y(X).shape
+    model.fit(X, y)
+    shape_after_fit = model.sample_y(X).shape
+    assert shape_before_fit == shape_after_fit
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_predict_shape_with_prior(n_targets):
+    """Check the output shape of `predict` with prior distribution"""
+    rng = np.random.RandomState(1024)
+
+    n_sample = 10
+    X = rng.randn(n_sample, 3)
+    y = rng.randn(n_sample, n_targets or 1)
+
+    model = GaussianProcessRegressor(n_targets=n_targets)
+    mean_prior, cov_prior = model.predict(X, return_cov=True)
+    _, std_prior = model.predict(X, return_std=True)
+
+    model.fit(X, y)
+    mean_post, cov_post = model.predict(X, return_cov=True)
+    _, std_post = model.predict(X, return_std=True)
+
+    assert mean_prior.shape == mean_post.shape
+    assert cov_prior.shape == cov_post.shape
+    assert std_prior.shape == std_post.shape


Code review comments:
diff_hunk: @@ -378,12 +386,26 @@ def predict(self, X, return_std=False, return_cov=False):
                 )
             else:
                 kernel = self.kernel
-            y_mean = np.zeros(X.shape[0])
+
+            n_targets = self.n_targets or 1
path: sklearn/gaussian_process/_gpr.py
text: @user1:
I find the following more explicit.

```suggestion
            n_targets = self.n_targets if self.n_targets is not None else 1
```

diff_hunk: @@ -764,3 +764,40 @@ def test_sample_y_shapes(normalize_y, n_targets):
 
     y_samples = model.sample_y(X_test, n_samples=n_samples_y_test)
     assert y_samples.shape == y_test_shape
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_sample_y_shape_with_prior(n_targets):
path: sklearn/gaussian_process/tests/test_gpr.py
text: @user1:
```suggestion
def test_sample_y_shape_with_prior(n_targets, n_samples):
```

diff_hunk: @@ -764,3 +764,40 @@ def test_sample_y_shapes(normalize_y, n_targets):
 
     y_samples = model.sample_y(X_test, n_samples=n_samples_y_test)
     assert y_samples.shape == y_test_shape
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_sample_y_shape_with_prior(n_targets):
+    """Check the output shape of `sample_y` is consistent before and after `fit`"""
+    rng = np.random.RandomState(1024)
+
+    X = rng.randn(10, 3)
+    y = rng.randn(10, n_targets or 1)
path: sklearn/gaussian_process/tests/test_gpr.py
text: @user1:
```suggestion
    y = rng.randn(10, n_targets if n_targets is not None else 1)
```

diff_hunk: @@ -108,6 +108,12 @@ def optimizer(obj_func, initial_theta, bounds):
         which might cause predictions to change if the data is modified
         externally.
 
+    n_targets : int, default=None
+        Dimension of the target values. Used to decide the output dimension
+        when sampling with prior distribution.
path: sklearn/gaussian_process/_gpr.py
text: @user1:
```suggestion
        The number of dimensions of the target values. Used to decide the number
        of outputs when sampling from the prior distributions (i.e. calling
        :meth:`sample_y` before :meth:`fit`). This parameter is ignored once
        :meth:`fit` has been called.
```

diff_hunk: @@ -764,3 +764,40 @@ def test_sample_y_shapes(normalize_y, n_targets):
 
     y_samples = model.sample_y(X_test, n_samples=n_samples_y_test)
     assert y_samples.shape == y_test_shape
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_sample_y_shape_with_prior(n_targets):
+    """Check the output shape of `sample_y` is consistent before and after `fit`"""
+    rng = np.random.RandomState(1024)
+
+    X = rng.randn(10, 3)
+    y = rng.randn(10, n_targets or 1)
+
+    model = GaussianProcessRegressor(n_targets=n_targets)
+    shape_before_fit = model.sample_y(X).shape
path: sklearn/gaussian_process/tests/test_gpr.py
text: @user1:
```suggestion
    shape_before_fit = model.sample_y(X, n_samples=n_samples).shape
```

diff_hunk: @@ -764,3 +764,40 @@ def test_sample_y_shapes(normalize_y, n_targets):
 
     y_samples = model.sample_y(X_test, n_samples=n_samples_y_test)
     assert y_samples.shape == y_test_shape
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_sample_y_shape_with_prior(n_targets):
+    """Check the output shape of `sample_y` is consistent before and after `fit`"""
+    rng = np.random.RandomState(1024)
+
+    X = rng.randn(10, 3)
+    y = rng.randn(10, n_targets or 1)
+
+    model = GaussianProcessRegressor(n_targets=n_targets)
+    shape_before_fit = model.sample_y(X).shape
+    model.fit(X, y)
+    shape_after_fit = model.sample_y(X).shape
path: sklearn/gaussian_process/tests/test_gpr.py
text: @user1:
```suggestion
    shape_after_fit = model.sample_y(X, n_samples=n_samples).shape
```

diff_hunk: @@ -108,6 +108,12 @@ def optimizer(obj_func, initial_theta, bounds):
         which might cause predictions to change if the data is modified
         externally.
 
+    n_targets : int, default=None
+        Dimension of the target values. Used to decide the output dimension
+        when sampling with prior distribution.
+
+        .. versionadded:: 1.1
path: sklearn/gaussian_process/_gpr.py
text: @user1:
```suggestion
        .. versionadded:: 1.2
```

diff_hunk: @@ -764,3 +764,40 @@ def test_sample_y_shapes(normalize_y, n_targets):
 
     y_samples = model.sample_y(X_test, n_samples=n_samples_y_test)
     assert y_samples.shape == y_test_shape
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_sample_y_shape_with_prior(n_targets):
+    """Check the output shape of `sample_y` is consistent before and after `fit`"""
+    rng = np.random.RandomState(1024)
+
+    X = rng.randn(10, 3)
+    y = rng.randn(10, n_targets or 1)
+
+    model = GaussianProcessRegressor(n_targets=n_targets)
+    shape_before_fit = model.sample_y(X).shape
+    model.fit(X, y)
+    shape_after_fit = model.sample_y(X).shape
+    assert shape_before_fit == shape_after_fit
+
+
+@pytest.mark.parametrize("n_targets", [None, 1, 2, 3])
+def test_predict_shape_with_prior(n_targets):
+    """Check the output shape of `predict` with prior distribution"""
+    rng = np.random.RandomState(1024)
+
+    n_sample = 10
+    X = rng.randn(n_sample, 3)
+    y = rng.randn(n_sample, n_targets or 1)
path: sklearn/gaussian_process/tests/test_gpr.py
text: @user1:
```suggestion
    y = rng.randn(n_sample, n_targets if n_targets is not None else 1)
```

diff_hunk: @@ -378,12 +386,26 @@ def predict(self, X, return_std=False, return_cov=False):
                 )
             else:
                 kernel = self.kernel
-            y_mean = np.zeros(X.shape[0])
+
+            n_targets = self.n_targets or 1
+            if n_targets == 1:
+                y_mean = np.zeros(X.shape[0])
+            else:
+                y_mean = np.zeros(shape=(X.shape[0], self.n_targets))
path: sklearn/gaussian_process/_gpr.py
text: @user1:
```suggestion
            y_mean = np.zeros(shape=(X.shape[0], self.n_targets)).squeeze()
```