<commit_before>
import logging
import george
import numpy as np

from scipy import optimize
from copy import deepcopy

from robo.models.base_model import BaseModel

logger = logging.getLogger(__name__)


class GaussianProcess(BaseModel):

    def __init__(self, kernel, prior=None,
                 noise=1e-3, use_gradients=False,
                 basis_func=None, dim=None, normalize_output=False,
                 *args, **kwargs):
        """
        Interface to the george GP library. The GP hyperparameter are obtained
        by optimizing the marginal loglikelihood.

        Parameters
        ----------
        kernel : george kernel object
            Specifies the kernel that is used for all Gaussian Process
        prior : prior object
            Defines a prior for the hyperparameters of the GP. Make sure that
            it implements the Prior interface.
        noise : float
            Noise term that is added to the diagonal of the covariance matrix
            for the cholesky decomposition.
        use_gradients : bool
            Use gradient information to optimize the negative log likelihood
        """

        self.kernel = kernel
        self.model = None
        self.prior = prior
        self.noise = noise
        self.use_gradients = use_gradients
        self.basis_func = basis_func
        self.dim = dim
        self.normalize_output = normalize_output
        self.X = None
        self.Y = None

    def scale(self, x, new_min, new_max, old_min, old_max):
        return ((new_max - new_min) *
                (x - old_min) / (old_max - old_min)) + new_min

    def train(self, X, Y, do_optimize=True):
        """
        Computes the cholesky decomposition of the covariance of X and
        estimates the GP hyperparameter by optimizing the marginal
        loglikelihood. The prior mean of the GP is set to the empirical
        mean of the X.

        Parameters
        ----------
        X: np.ndarray (N, D)
            Input data points. The dimensionality of X is (N, D),
            with N as the number of points and D is the number of features.
        Y: np.ndarray (N, 1)
            The corresponding target values.
        do_optimize: boolean
            If set to true the hyperparameters are optimized.
        """
        self.X = X
        # For Fabolas we transform s to (1 - s)^2
        if self.basis_func is not None:
            self.X = deepcopy(X)
            self.X[:, self.dim] = self.basis_func(self.X[:, self.dim])

        self.Y = Y
        if self.normalize_output:
            self.Y_mean = np.mean(Y)
            self.Y_std = np.std(Y)
            self.Y = (Y - self.Y_mean) / self.Y_std

        # Use the empirical mean of the data as mean for the GP
        self.mean = np.mean(self.Y, axis=0)
        self.model = george.GP(self.kernel, mean=self.mean)

        if do_optimize:
            self.hypers = self.optimize()
            self.model.kernel[:] = self.hypers[:-1]
            self.noise = np.exp(self.hypers[-1]) ## sigma^2
        else:
            self.hypers = self.model.kernel[:]
            self.hypers = np.append(self.hypers, np.log(self.noise))
        logger.debug("HYPERS: " + str(self.hypers))
        self.model.compute(self.X, yerr=np.sqrt(self.noise))

    def get_noise(self):
        return self.noise

    def nll(self, theta):
        """
        Returns the negative marginal log likelihood (+ the prior) for
        a hyperparameter configuration theta.
        (negative because we use scipy minimize for optimization)

        Parameters
        ----------
        theta : np.ndarray(H)
            Hyperparameter vector. Note that all hyperparameter are
            on a log scale.

        Returns
        ----------
        float
            lnlikelihood + prior
        """
        # Specify bounds to keep things sane
        if np.any((-20 > theta) + (theta > 20)):
            return 1e25

        # The last entry of theta is always the noise
        self.model.kernel[:] = theta[:-1]
        noise = np.exp(theta[-1])  # sigma^2
        
        self.model.compute(self.X, yerr=np.sqrt(noise))
        ll = self.model.lnlikelihood(self.Y[:, 0], quiet=True)

        # Add prior
        if self.prior is not None:
            ll += self.prior.lnprob(theta)

        # We add a minus here because scipy is minimizing
        return -ll if np.isfinite(ll) else 1e25

    def grad_nll(self, theta):

        self.model.kernel[:] = theta[:-1]
        noise = np.exp(theta[-1])
        
        self.model.compute(self.X, yerr=np.sqrt(noise))
        
        self.model._compute_alpha(self.Y[:, 0])
        K_inv = self.model.solver.apply_inverse(np.eye(self.model._alpha.size),
                                          in_place=True)

        # The gradients of the Gram matrix, for the noise this is just 
        # the identiy matrix
        Kg = self.model.kernel.gradient(self.model._x)
        Kg = np.concatenate((Kg, np.eye(Kg.shape[0])[:, :, None]), axis=2)

        # Calculate the gradient.
        A = np.outer(self.model._alpha, self.model._alpha) - K_inv
        g = 0.5 * np.einsum('ijk,ij', Kg, A)
        

        if self.prior is not None:
            g += self.prior.gradient(theta)

        return -g

    def optimize(self):
        # Start optimization  from the previous hyperparameter configuration
        p0 = self.model.kernel.vector
        p0 = np.append(p0, np.log(self.noise))

        if self.use_gradients:
            bounds = [(-10, 10)] * (len(self.kernel) + 1)
            theta, _, _ = optimize.fmin_l_bfgs_b(self.nll, p0,
                                             fprime=self.grad_nll,
                                             bounds=bounds)
        else:
            results = optimize.minimize(self.nll, p0)
            theta = results.x
        return theta

    def predict_variance(self, X1, X2):
        r"""
        Predicts the variance between two test points X1, X2 by
           math: \sigma(X_1, X_2) = k_{X_1,X_2} - k_{X_1,X} * (K_{X,X}
                       + \sigma^2*\mathds{I})^-1 * k_{X,X_2})

        Parameters
        ----------
        X1: np.ndarray (N, D)
            First test point
        X2: np.ndarray (N, D)
            Second test point<commit_msg>pep8 fixes<commit_after>
import logging
import george
import numpy as np

from scipy import optimize
from copy import deepcopy

from robo.models.base_model import BaseModel

logger = logging.getLogger(__name__)


class GaussianProcess(BaseModel):

    def __init__(self, kernel, prior=None,
                 noise=1e-3, use_gradients=False,
                 basis_func=None, dim=None, normalize_output=False,
                 *args, **kwargs):
        """
        Interface to the george GP library. The GP hyperparameter are obtained
        by optimizing the marginal loglikelihood.

        Parameters
        ----------
        kernel : george kernel object
            Specifies the kernel that is used for all Gaussian Process
        prior : prior object
            Defines a prior for the hyperparameters of the GP. Make sure that
            it implements the Prior interface.
        noise : float
            Noise term that is added to the diagonal of the covariance matrix
            for the cholesky decomposition.
        use_gradients : bool
            Use gradient information to optimize the negative log likelihood
        """

        self.kernel = kernel
        self.model = None
        self.prior = prior
        self.noise = noise
        self.use_gradients = use_gradients
        self.basis_func = basis_func
        self.dim = dim
        self.normalize_output = normalize_output
        self.X = None
        self.Y = None
        self.hypers = []

    def scale(self, x, new_min, new_max, old_min, old_max):
        return ((new_max - new_min) *
                (x - old_min) / (old_max - old_min)) + new_min

    def train(self, X, Y, do_optimize=True):
        """
        Computes the cholesky decomposition of the covariance of X and
        estimates the GP hyperparameter by optimizing the marginal
        loglikelihood. The prior mean of the GP is set to the empirical
        mean of the X.

        Parameters
        ----------
        X: np.ndarray (N, D)
            Input data points. The dimensionality of X is (N, D),
            with N as the number of points and D is the number of features.
        Y: np.ndarray (N, 1)
            The corresponding target values.
        do_optimize: boolean
            If set to true the hyperparameters are optimized.
        """

        self.X = X

        # For Fabolas we transform s to (1 - s)^2
        if self.basis_func is not None:
            self.X = deepcopy(X)
            self.X[:, self.dim] = self.basis_func(self.X[:, self.dim])

        self.Y = Y
        if self.normalize_output:
            self.Y_mean = np.mean(Y)
            self.Y_std = np.std(Y)
            self.Y = (Y - self.Y_mean) / self.Y_std

        # Use the empirical mean of the data as mean for the GP
        self.mean = np.mean(self.Y, axis=0)
        self.model = george.GP(self.kernel, mean=self.mean)

        if do_optimize:
            self.hypers = self.optimize()
            self.model.kernel[:] = self.hypers[:-1]
            self.noise = np.exp(self.hypers[-1]) ## sigma^2
        else:
            self.hypers = self.model.kernel[:]
            self.hypers = np.append(self.hypers, np.log(self.noise))
        logger.debug("HYPERS: " + str(self.hypers))
        self.model.compute(self.X, yerr=np.sqrt(self.noise))

    def get_noise(self):
        return self.noise

    def nll(self, theta):
        """
        Returns the negative marginal log likelihood (+ the prior) for
        a hyperparameter configuration theta.
        (negative because we use scipy minimize for optimization)

        Parameters
        ----------
        theta : np.ndarray(H)
            Hyperparameter vector. Note that all hyperparameter are
            on a log scale.

        Returns
        ----------
        float
            lnlikelihood + prior
        """
        # Specify bounds to keep things sane
        if np.any((-20 > theta) + (theta > 20)):
            return 1e25

        # The last entry of theta is always the noise
        self.model.kernel[:] = theta[:-1]
        noise = np.exp(theta[-1])  # sigma^2
        
        self.model.compute(self.X, yerr=np.sqrt(noise))
        ll = self.model.lnlikelihood(self.Y[:, 0], quiet=True)

        # Add prior
        if self.prior is not None:
            ll += self.prior.lnprob(theta)

        # We add a minus here because scipy is minimizing
        return -ll if np.isfinite(ll) else 1e25

    def grad_nll(self, theta):

        self.model.kernel[:] = theta[:-1]
        noise = np.exp(theta[-1])
        
        self.model.compute(self.X, yerr=np.sqrt(noise))
        
        self.model._compute_alpha(self.Y[:, 0])
        K_inv = self.model.solver.apply_inverse(np.eye(self.model._alpha.size),
                                          in_place=True)

        # The gradients of the Gram matrix, for the noise this is just 
        # the identiy matrix
        Kg = self.model.kernel.gradient(self.model._x)
        Kg = np.concatenate((Kg, np.eye(Kg.shape[0])[:, :, None]), axis=2)

        # Calculate the gradient.
        A = np.outer(self.model._alpha, self.model._alpha) - K_inv
        g = 0.5 * np.einsum('ijk,ij', Kg, A)

        if self.prior is not None:
            g += self.prior.gradient(theta)

        return -g

    def optimize(self):
        # Start optimization from the previous hyperparameter configuration
        p0 = self.model.kernel.vector
        p0 = np.append(p0, np.log(self.noise))

        if self.use_gradients:
            bounds = [(-10, 10)] * (len(self.kernel) + 1)
            theta, _, _ = optimize.fmin_l_bfgs_b(self.nll, p0,
                                             fprime=self.grad_nll,
                                             bounds=bounds)
        else:
            results = optimize.minimize(self.nll, p0)
            theta = results.x
        return theta

    def predict_variance(self, X1, X2):
        r"""
        Predicts the variance between two test points X1, X2 by
           math: \sigma(X_1, X_2) = k_{X_1,X_2} - k_{X_1,X} * (K_{X,X}
                       + \sigma^2*\mathds{I})^-1 * k_{X,X_2})

        Parameters
        ----------
        X1: np.ndarray (N, D)
            First test point
        X2: np.ndarray (N, D)
            Second test point
{'fraction_non_alphanumeric': 0.10153256704980843, 'fraction_numerical': 0.008700510855683269, 'mean_word_length': 1.6978897502153316, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '904361', 'n_tokens_mistral': 3813, 'n_tokens_neox': 3596, 'n_words': 1079}