Gradient of a Gaussian process regressor from scikit-learn

gaussian_process_regressor_gradient.py

import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor


def gradient(gpr: GaussianProcessRegressor, Xstar: np.ndarray) -> np.ndarray:
    """Evaluate the gradient of a Gaussian process at a given point."""
    X = gpr.X_train_
    m, n = Xstar.shape[0], X.shape[0]
    assert Xstar.shape[1] == X.shape[1]
    Xstar_minus_X = Xstar[:, np.newaxis, :] - X[np.newaxis, :, :]
    kXstarX = gpr.kernel_(Xstar, X)
    length_scale = gpr.kernel_.get_params()['k2__length_scale']
    kXstarX_times_alpha = (kXstarX * gpr.alpha_)[:, :, np.newaxis]
    return -(kXstarX_times_alpha * Xstar_minus_X).sum(axis=1) / length_scale**2 * gpr._y_train_std