diff --git a/autoemulate/emulators/__init__.py b/autoemulate/emulators/__init__.py index 9b46a9a0..b0ff3184 100644 --- a/autoemulate/emulators/__init__.py +++ b/autoemulate/emulators/__init__.py @@ -3,10 +3,12 @@ from .gaussian_process_sk import GaussianProcessSk from .neural_network import NeuralNetwork from .random_forest import RandomForest +from .radial_basis import RadialBasis MODEL_REGISTRY = { # "GaussianProcess": GaussianProcess, "GaussianProcessSk": GaussianProcessSk, "NeuralNetwork": NeuralNetwork, "RandomForest": RandomForest, + "RadialBasis": RadialBasis, } diff --git a/autoemulate/emulators/radial_basis.py b/autoemulate/emulators/radial_basis.py index b28a5ca9..43a8b441 100644 --- a/autoemulate/emulators/radial_basis.py +++ b/autoemulate/emulators/radial_basis.py @@ -38,14 +38,17 @@ def fit(self, X, y): self : object Returns self. """ - X, y = check_X_y(X, y, multi_output=True, y_numeric=True) + X, y = check_X_y(X, y, multi_output=True, y_numeric=True, dtype=np.float64) self.n_features_in_ = X.shape[1] - self.model_ = RBF(d0=self.d0, poly_degree=self.poly_degree, reg=self.reg) + self.model_ = RBF( + d0=self.d0, poly_degree=self.poly_degree, reg=self.reg, print_global=False + ) self.model_.set_training_values(X, y) self.model_.train() + self.is_fitted_ = True return self - def predict(self, X, return_std=False): + def predict(self, X): """Predicts the output of the emulator for a given input. Parameters @@ -53,28 +56,28 @@ def predict(self, X, return_std=False): X : {array-like, sparse matrix}, shape (n_samples, n_features) The input samples. - return_std : bool, default=False - Whether to return the standard deviation of the prediction. - Returns ------- y : ndarray of shape (n_samples, n_features) The predicted values. """ - X = check_array(X) + X = check_array(X, dtype=np.float64) check_is_fitted(self) - if return_std: - return self.model_.predict_values(X), self.model_.predict_variances(X) - else: - return self.model_.predict_values(X) + predictions = self.model_.predict_values(X) + # Check if dimension 2 is 1, if so, reshape to 1D array + # This is needed to pass sklearn's check_estimator + if predictions.ndim == 2 and predictions.shape[1] == 1: + predictions = predictions.ravel() + return predictions def get_grid_params(self): """Returns the grid parameters of the emulator.""" param_grid = { - "d0": [1.0, 2.0, 3.0], + "d0": [1.0], "poly_degree": [-1, 0, 1], "reg": [1e-10, 1e-5, 1e-2], } + return param_grid def _more_tags(self): return {"multioutput": True} diff --git a/tests/test_estimators.py b/tests/test_estimators.py index 63970701..22a5326e 100644 --- a/tests/test_estimators.py +++ b/tests/test_estimators.py @@ -4,6 +4,7 @@ GaussianProcessSk, NeuralNetwork, GaussianProcess, + RadialBasis, ) from functools import partial @@ -13,6 +14,7 @@ RandomForest(random_state=42), GaussianProcessSk(random_state=1337), NeuralNetwork(random_state=13), + # RadialBasis(), ] ) def test_check_estimator(estimator, check):