From b78755718027b7874e1f96da3478940bb801c605 Mon Sep 17 00:00:00 2001 From: Mirochill <200482516+Mirochill@users.noreply.github.com> Date: Wed, 27 May 2026 11:03:18 +0200 Subject: [PATCH] Return GP predictions as column vectors --- pySOT/surrogate/gp.py | 2 +- tests/test_surrogates.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pySOT/surrogate/gp.py b/pySOT/surrogate/gp.py index 305bbb6..4269608 100644 --- a/pySOT/surrogate/gp.py +++ b/pySOT/surrogate/gp.py @@ -68,7 +68,7 @@ def predict(self, xx): """ self._fit() xx = to_unit_box(np.atleast_2d(xx), self.lb, self.ub) - return self._mu + self._sigma * self.model.predict(xx) + return self._mu + self._sigma * np.expand_dims(self.model.predict(xx), axis=1) def predict_std(self, xx): """Predict standard deviation at points xx. diff --git a/tests/test_surrogates.py b/tests/test_surrogates.py index 3cfbf97..d79dc05 100644 --- a/tests/test_surrogates.py +++ b/tests/test_surrogates.py @@ -174,6 +174,9 @@ def test_gp(): np.random.seed(0) Xs = np.random.rand(10, 2) fhx = gp.predict(Xs) + assert fhx.shape == (Xs.shape[0], 1) + assert gp.predict_std(Xs).shape == fhx.shape + assert gp.predict(Xs[0]).shape == (1, 1) fx = f(Xs) assert np.max(np.abs(fx - fhx)) < 1e-2