-
Notifications
You must be signed in to change notification settings - Fork 127
/
sklearn_model_wrapper.py
53 lines (41 loc) · 1.53 KB
/
sklearn_model_wrapper.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from typing import Tuple
import numpy as np
try:
from sklearn.gaussian_process import GaussianProcessRegressor
except ImportError:
ImportError("scikit-learn needs to be installed in order to use SklearnGPRWrapper")
from emukit.core.interfaces.models import IModel
class SklearnGPRWrapper(IModel):
def __init__(self, sklearn_model: GaussianProcessRegressor):
"""
:param sklearn_model: Scikit-learn GPR model to wrap
"""
self.model = sklearn_model
def predict(self, X: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
"""
Predict mean and variance values for given points
:param X: array of shape (n_points x n_inputs) of points to run prediction for
:return: Tuple of mean and variance which are 2d arrays of shape (n_points x n_outputs)
"""
mean, std = self.model.predict(X, return_std=True)
if mean.ndim == 1:
mean = mean[:, None]
return mean, np.power(std, 2.0).reshape(-1, 1)
def set_data(self, X: np.ndarray, Y: np.ndarray) -> None:
"""
Sets training data in model
:param X: new points
:param Y: function values at new points X
"""
self.model.X_train_, self.model.y_train_ = X, Y
def optimize(self) -> None:
"""
Optimize hyper-parameters of model
"""
self.model.fit(self.X, self.Y)
@property
def X(self) -> np.ndarray:
return self.model.X_train_
@property
def Y(self) -> np.ndarray:
return self.model.y_train_