sklearn wrapper for SparseGPRegressor

[igi-chris]

I recently implemented a scikit learn wrapper for the SparseGPRegressor (with help from Dan Cornford) and thought I would share it here in case it is useful for anyone else.

We wanted to keep this within the sklearn ecosystem because we use the Pipeline, various preprocessing options etc. However, the sklearn GPR has some limitations, most notably that there is no sparse option, so it is not very scalable.

Here is the wrapper code:

import warnings

from GPy.models import SparseGPRegression
from GPy.kern import Kern, Matern52, Matern32, Exponential, RBF, RatQuad, Linear
from GPy.kern.src.kern import CombinationKernel
from GPy.likelihoods import Gaussian
from GPy.likelihoods.link_functions import Identity
import numpy as np
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.dummy import check_random_state
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted


# guidance on how to implement a custom sklearn regressor:
# https://scikit-learn.org/stable/developers/develop.html#rolling-your-own-estimator
class GPySparseRegressor(BaseEstimator, RegressorMixin):
    """
    A Gaussian Process (GP) regressor wrapper for GPy, compatible with scikit-learn.
    """

    # ⚠️ avoid using GPy object directly here e.g. GPy.kern as we need custom serialisation
    def __init__(self, random_state:int|None=None, 
                 num_inducing: int = 50, 
                 kernel: str = "Matern52", 
                 initial_kernel_variance: float = 1, 
                 initial_kernel_lengthscale: float = 1,
                 add_linear_kernel: bool = True,
                 linear_variance: float = 1, 
                 max_iter: int = 1000,
                 optimizer: str = 'lbfgs'):
        """
        Initializes the GP regressor.
        
        Parameters:
        random_state: int
            Allows for deterministic  selection of inducing points.
        optimizer: str, optional
            The optimizer to use for optimizing GP hyperparameters.
        details for the other parameters are captured in GPyRegressorParameters.
        """
        self.num_inducing = num_inducing
        self.kernel = kernel
        self.initial_kernel_variance = initial_kernel_variance
        self.initial_kernel_lengthscale = initial_kernel_lengthscale
        self.optimizer = optimizer
        self.max_iter = max_iter
        self.random_state = random_state
        self.add_linear_kernel = add_linear_kernel
        self.linear_variance = linear_variance

    def __str__(self):
        # overriding and excluding model so that we don't get "object at 0x00..."
        # in the model description
        incl_fields = {k: v for k, v in self.__dict__.items() if k != 'model'}
        fields_str = ', '.join(f"{k}={v}" for k, v in incl_fields.items())
        return f"{self.__class__.__name__}({fields_str})"

    def __repr__(self):
        return self.__str__()

    def fit(self, X, y):
        """
        Fit the Gaussian Process model to the training data.

        Parameters:
        X: array-like of shape (n_samples, n_features)
            Training data.
        y: array-like, shape (n_samples,) or (n_samples, n_targets)
            Target values.

        Returns:
        self: returns an instance of self.
        """
        # Validate that X and y are correctly formatted, non-empty, and have consistent dimensions
        X, y = check_X_y(X, y, y_numeric=True, estimator=self)
        self.n_features_in_ = X.shape[1]
        
        kernel = self._create_kernel(input_dim=X.shape[1])

        # Reshape y to be 2-dimensional if it's not already
        if hasattr(y, 'values'): # e.g. if y is passed in as a pandas series
            y = y.values
        y_2d = y.reshape(-1, 1) if y.ndim == 1 else y

        Z = self._generate_inducing_points(X)
        self.model_ = SparseGPRegression(X, y_2d, kernel=kernel, Z=Z)
        self.model_.optimize(self.optimizer, max_iters=self.max_iter)

        self._set_optimised_parameters()
        self.n_iter_ = self.max_iter  # Set to max_iter as an estimate (GPy doesn't expose this)
        self.is_fitted_ = True
        return self

    def predict(self, X, return_std:bool=False, with_obs_noise:bool=False):
        """
        Predict using the Gaussian Process model.

        Parameters:
        X: array-like of shape (n_samples, n_features)
            Samples to predict.
        return_std: bool, optional
            Whether to return standard deviation of predictions.
        with_obs_noise:
            Whether to include the observation noise variance in the prediction.
            Note: This does not change the mean, just the variance. We recommend
            including observation noise when assessing fit, not when using for 
            predictions. We assume the noise is independent and identically
            distributed (iid). When predicting with new data we do not know if the
            observation uncertainty for the new data is similar to the training data
            e.g. different instruments could be used. We want predictions to reflect 
            belief about the underlying system, not the observation noise.

        Returns:
        y_mean: array of shape (n_samples,)
            Mean of the Gaussian process distribution at the query points.
        y_std: array of shape (n_samples,), optional
            Standard deviation of the Gaussian process distribution at the query points.
        """
        check_is_fitted(self, 'model_') # looks for a `model_` attribute with a value

        X = check_array(X, estimator=self)

        if with_obs_noise:
            y_mean, y_var = self.model_.predict(X)
        else:
            lik = Gaussian(gp_link=Identity(), variance=1e-9)
            y_mean, y_var = self.model_.predict(X, include_likelihood=True,
                                               likelihood=lik)

        if return_std:
            y_std = np.sqrt(y_var)
            return y_mean.ravel(), y_std.ravel()

        return y_mean.ravel()

    def fit_predict(self, X, y, return_std:bool=False, with_obs_noise:bool=False):
        """
        Fit the Gaussian Process model to the training data and make predictions.

        Parameters:
        X: array-like of shape (n_samples, n_features)
            Training data.
        y: array-like, shape (n_samples,) or (n_samples, n_targets)
            Target values.
        return_std: bool, optional
            Whether to return standard deviation of predictions.
        with_obs_noise:
            Whether to include the observation noise variance in the prediction.
            y=f(x)+e. The mean does not change, but the variance does. Include
            noise when assessing fit, not when using for predictions.

        Returns:
        y_mean: array of shape (n_samples,)
            Mean of the Gaussian process distribution at the query points.
        y_std: array of shape (n_samples,), optional
            Standard deviation of the Gaussian process distribution at the query points.
        """
        
        self.fit(X, y)
        return self.predict(X, return_std=return_std, with_obs_noise=with_obs_noise)
    
    def _create_kernel(self, input_dim: int) -> Kern:
        # Primary kernel
        kernel =  self._create_single_kernel(self.kernel, input_dim)
    
        if self.add_linear_kernel:
            linear = self._create_single_kernel('Linear', input_dim,
                                                variance=self.linear_variance)
            kernel += linear
        return kernel
        
    def _create_single_kernel(self, kernel_type: str, input_dim: int,
                              variance:float|None=None) -> Kern:
        variance = self.initial_kernel_variance if variance is None else variance
        lengthscale = self.initial_kernel_lengthscale
        match kernel_type:
            case 'RBF':
                return RBF(input_dim, variance=variance, lengthscale=lengthscale)
            case 'Matern52':
                return Matern52(input_dim, variance=variance, lengthscale=lengthscale)
            case 'Matern32':
                return Matern32(input_dim, variance=variance, lengthscale=lengthscale)
            case 'RQ':
                return RatQuad(input_dim, variance=variance, lengthscale=lengthscale)
            case 'Exp':
                return Exponential(input_dim, variance=variance, lengthscale=lengthscale)
            case 'Linear':
                return Linear(input_dim, variances=variance)
            case _:
                raise ValueError(f"Unsupported kernel type: {kernel_type}")

    def _generate_inducing_points(self, X):
        """
        Generate inducing points for SparseGPRegressor in a deterministic way.

        Parameters:
        X: array-like of shape (n_samples, n_features)
            Training data.

        Returns:
        np.ndarray: Inducing points of shape (num_inducing, n_features).
        """
        # We generate inducing points rather than using the num_inducing arg in the GPy 
        # model class as this allows us to use a random seed to make this reproducible.
        num_data_rows = X.shape[0]

        self.num_inducing_ = self.num_inducing
        if self.num_inducing_ > num_data_rows:
            self.num_inducing_ = num_data_rows

        if self.num_inducing_ == num_data_rows:
            return X
        
        # Convert seed to numpy RandomState object
        self.random_state_ = check_random_state(self.random_state)
        inducing_point_indices = self.random_state_.choice(num_data_rows, 
                                                           size=self.num_inducing_, 
                                                           replace=False)

        # Select the data points for random indices as inducing points
        inducing_points = X[inducing_point_indices, :]

        return inducing_points
    
    def _set_optimised_parameters(self):
        self.kernel_variance_ = None
        self.kernel_lengthscale_ = None
        self.linear_variance_ = None
        self.gaussian_noise_variance_ = None

        try:
            kern = self.model_.kern
            k = kern.parts[0] if isinstance(kern, CombinationKernel) else kern

            if hasattr(k, 'variance'):
                self.kernel_variance_ = k.variance.values[0]

            if hasattr(k, 'lengthscale'):
                self.kernel_lengthscale_ = k.lengthscale.values[0]

            if hasattr(k, 'variances'):  # linear kernel
                self.linear_variance_ = k.variances.values[0]

            if isinstance(kern, CombinationKernel) and len(kern.parts) > 1:
                lk = kern.parts[1]
                self.linear_variance_ = lk.variances.values[0]

            if hasattr(self.model_, 'Gaussian_noise'):
                self.gaussian_noise_variance_ = self.model_.Gaussian_noise.variance.values[0]

        except Exception as e:
            warnings.warn(f"Unable to set optimised parameters from kernel {k.name}: {e}")

We have used sklearn's check_estimator utility to validate this:

from unittest import TestCase
from sklearn.utils.estimator_checks import check_estimator

class TestGPyWrapper(TestCase):

    def test_it_is_a_valid_sklearn_estimator(self):
        model = GPySparseRegressor()
        # use an sklearn util to check that the wrapper conforms to sklearn conventions
        # this will raise an error and fail if the model does not pass the checks
        check_estimator(model)