Source code for ariane.lib.gaussian_process_regression.implementations.sklearn._sklearn

from ... import GaussianProcessRegressorImplementation
from ....utils.math import cholesky_update

import numpy as np
from scipy.linalg import cho_solve
from sklearn.gaussian_process import GaussianProcessRegressor as sklearn_GPR
from sklearn.exceptions import ConvergenceWarning
import warnings

warnings.filterwarnings("ignore", category=ConvergenceWarning)




class ScikitLearnGaussianProcessRegressor(GaussianProcessRegressorImplementation):
    """Wrapper class for scikit-learn's implementation of a Gaussian process regressor."""

    def __init__(self, kernel=None, *, noise=1e-10,
                 optimizer="fmin_l_bfgs_b", n_restarts_optimizer=0,
                 normalize_y=False, copy_X_train=True, random_state=None):
        self.gpr = sklearn_GPR(kernel=kernel, alpha=noise, optimizer=optimizer,
                               n_restarts_optimizer=n_restarts_optimizer,
                               normalize_y=normalize_y, copy_X_train=copy_X_train,
                               random_state=random_state)

    def __call__(self, x, return_std=False):
        return self.gpr.predict(x, return_std=return_std)



    def fit(self, x, y, noise=None):
        if noise is not None:
            self.noise = noise

        self.gpr.fit(x, y)


    @property
    def x_train(self):
        return self.gpr.X_train_

    @property
    def y_train(self):
        return self.gpr.y_train_

    @property
    def noise(self):
        return self.gpr.alpha

    @noise.setter
    def noise(self, noise):
        self.gpr.alpha = noise

    @property
    def kernel(self):
        return self.gpr.kernel_

    @property
    def kernel_hyperparameters(self):
        return np.exp(self.gpr.kernel_.theta)

    @property
    def kernel_length_scales(self):
        return self.kernel.k2.length_scale



    def add_data(self, x_new, y_new, noise=None):
        x_new = np.asarray(x_new)
        y_new = np.asarray(y_new)

        # check noise compatibility
        if np.isscalar(self.noise):
            if noise is None:
                noise = self.noise
            else:
                raise Exception('scalar noise was already specified at initialization')
        else:
            if np.iterable(noise):
                assert len(noise) == len(y_new)
                noise = np.asarray(noise)
                self.noise = np.append(self.noise, noise)
            else:
                raise Exception('noise needs to be iterable')

        x_train_new = np.append(self.x_train, x_new, axis=0)
        y_train_new = np.append(self.y_train, y_new)

        if self.gpr.optimizer is not None:
            # make new optimization of kernel hyperparameters through `fit`
            self.fit(x_train_new, y_train_new)
        else:
            # update existing Cholesky decomposition (`self.gpr.L_`, `self.gpr.alpha_`)
            Kij_new = self.kernel(self.x_train, x_new)
            Kii_new = self.kernel(x_new, x_new)
            Kii_new[np.diag_indices_from(Kii_new)] += noise

            self.gpr.X_train_ = x_train_new
            self.gpr.y_train_ = y_train_new

            # update
            self.gpr.L_ = cholesky_update(self.gpr.L_, Kij_new, Kii_new)
            self.gpr.alpha_ = cho_solve((self.gpr.L_, True), self.y_train)

            # `self.gpr.L_` changed, `self.gpr._K_inv` needs to be recomputed for predictive variances
            self.gpr._K_inv = None




    def stop_optimization(self):
        self.gpr.kernel = self.gpr.kernel_
        self.gpr.optimizer = None