Kernels¶

Contains all internal kernels for the GPR. New Kernels can be created by inheriting from the base kernel ‘Kernel’

Example

>>> from ife_surrogate.gp.kernels import Kernel
>>> from flax import struct
>>>
>>> @struct.dataclass
>>> class NewKernel(Kernel):
>>>     def evaluate(self, x1: Array, x2: Array) -> Array:
>>>         # some kernel definition
>>>         return []

class ife_surrogate.gp.kernels.Kernel(*, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: ABC

Abstract base class for kernels in Gaussian Process (GP) regression.

This class provides the core interface and utility functions for all kernel implementations, including methods to evaluate the kernel, access and update parameters, and sample new hyperparameters from prior distributions.

Parameters:

priors (Dict[str, Callable])
param_bounds (Dict)
exempt (Set[str])

priors¶

A dictionary mapping parameter names to callable prior distributions used for sampling hyperparameters.

Type:: Dict[str, Callable]

param_bounds¶

Dictionary specifying bounds for each parameter. Defaults to an empty dictionary.

Type:: Dict, optional

exempt¶

Set of attribute names to exclude when retrieving parameters (e.g., priors, param_bounds, exempt itself).

Type:: Set[str]

evaluate(x1, x2)¶

Abstract method. Computes the kernel function between two inputs.

Parameters:

x1 (Array)
x2 (Array)

Return type:

Array

__call__(x1, x2)¶: Calls evaluate, allowing the kernel object to be used as a function.

get_priors()¶: Returns the dictionary of priors for kernel hyperparameters.

get_params()¶

Returns a dictionary of the current kernel parameters, excluding those in exempt.

Return type:: Dict[str, Array]

sample_hyperparameters(key)¶

Samples new parameter values from the prior distributions using a JAX random key.

Parameters:: key (Key)
Return type:: Dict[str, Array]

update_params(params)¶

Updates kernel parameters in-place using the provided dictionary.

Parameters:: params (Dict[str, Array])
Return type:: None

priors: Dict[str, Callable]¶

param_bounds: Dict¶

exempt: Set[str]¶

abstract evaluate(x1, x2)¶

Evaluate the kernel function.

Parameters:

x1 (Array) – First input data array.
x2 (Array) – Second input data array.

Returns:

Kernel evaluation between x1 and x2.

Return type:

Array

get_priors()¶

Get the kernel parameters.

Returns:: A dictionary containing the kernel parameters.

get_params()¶

Get the kernel parameters.

Returns:: A dictionary containing the kernel parameters.
Return type:: Dict[str, Array]

sample_hyperparameters(key)¶

Sample new parameters from the prior distribution.

Parameters:: key (Key) – JAX random key.
Returns:: A dictionary containing sampled parameters.
Return type:: Dict[str, Array]

update_params(params)¶

Update the kernel parameters.

Parameters:: params (Dict[str, Array]) – A dictionary containing the new parameters.
Return type:: None

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.RBF(lengthscale, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Radial Basis Function (RBF) Kernel:

\[k(\mathbf{x}_1, \mathbf{x}_2) = \exp\left( -\frac{\lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert^2}{2\ell^2} \right)\]

Parameters:

lengthscale (Array) – Lengthscale controlling the smoothness of the kernel.
priors (Dict[str, Callable], optional) – Priors for hyperparameters.
param_bounds (Dict)
exempt (Set[str])

lengthscale: Array¶

evaluate(x1, x2)¶

Evaluate the Radial Basis Function (RBF) kernel between x1 and x2.

Parameters:

x1 (Array) – Input array of shape (n_samples_1, n_features).
x2 (Array) – Input array of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2).

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.Kriging(lengthscale, power, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Kriging kernel (a generalized RBF kernel with a power term).

Computes the covariance between two inputs x1 and x2 as:

\[k_x(\mathbf x, \mathbf{x'}; \mathbf \theta) = \exp{\left(-\frac{1}{2}\sum_{i=1}^d \frac{1}{l_i^{2}} |x_{i} - x'_{i}|^{\kappa_i}\right)}\]

where: - lengthscale controls the smoothness along each input dimension. - power exponent of the absolute distance .

Parameters:

lengthscale (Array) – Positive scaling factors for each input dimension.
power (Array) – Power applied to the distance; must be > 0.
priors (Dict[str, Callable], optional) – Optional priors for hyperparameters.
param_bounds (Dict)
exempt (Set[str])

lengthscale: Array¶

power: Array¶

evaluate(x1, x2)¶

Evaluate the Kriging kernel between x1 and x2.

Parameters:

x1 (Array) – Inputs of shape (n_samples_1, n_features).
x2 (Array) – Inputs of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2).

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.RQ(lengthscale, alpha, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Rational Quadratic Kernel.

The covariance function is given by:

\[k(\mathbf{x}_1, \mathbf{x}_2) = \left(1 + \frac{\lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert^2}{2\alpha \ell^2} \right)^{-\alpha}\]

Parameters:

lengthscale (Array) – Lengthscale controlling the smoothness of the kernel.
alpha (float) – The alpha parameter, controls the relative scale of large and small lengthscales.
priors (Dict[str, Callable], optional) – Priors for hyperparameters.
param_bounds (Dict)
exempt (Set[str])

lengthscale: Array¶

alpha: float¶

evaluate(x1, x2)¶

Evaluate the Rational Quadratic kernel between x1 and x2.

Parameters:

x1 (Array) – Input array of shape (n_samples_1, n_features).
x2 (Array) – Input array of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2).

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.Matern12(lengthscale, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Matérn 1/2 kernel (Exponential kernel).

Computes the covariance between x1 and x2 as:

\[k(\mathbf{x}_1, \mathbf{x}_2) = \exp\left(-\frac{\lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert}{\ell}\right)\]

Parameters:

lengthscale (Array) – Lengthscale controlling the smoothness.
priors (Dict[str, Callable], optional) – Priors for hyperparameters.
param_bounds (Dict)
exempt (Set[str])

lengthscale: Array¶

evaluate(x1, x2)¶

Evaluate the Matérn 1/2 kernel.

Parameters:

x1 (Array) – Input array of shape (n_samples_1, n_features).
x2 (Array) – Input array of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2).

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.Matern32(lengthscale, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Matérn 3/2 kernel.

Computes the covariance between x1 and x2 as:

\[k(\mathbf{x}_1, \mathbf{x}_2) = \left(1 + \frac{\sqrt{3} \lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert}{\ell} \right) \exp\left(-\frac{\sqrt{3} \lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert}{\ell} \right)\]

Parameters:

lengthscale (Array) – Lengthscale controlling the smoothness.
priors (Dict[str, Callable], optional) – Priors for hyperparameters.
param_bounds (Dict)
exempt (Set[str])

lengthscale: Array¶

evaluate(x1, x2)¶

Evaluate the Matérn 3/2 kernel.

Parameters:

x1 (Array) – Input array of shape (n_samples_1, n_features).
x2 (Array) – Input array of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2).

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.Matern52(lengthscale, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Matérn 5/2 kernel.

Computes the covariance between x1 and x2 as:

\[k(\mathbf{x}_1, \mathbf{x}_2) = \left(1 + \frac{\sqrt{5} \lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert}{\ell} + \frac{5 \lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert^2}{3 \ell^2} \right) \exp\left(-\frac{\sqrt{5} \lVert \mathbf{x}_1 - \mathbf{x}_2 \rVert}{\ell} \right)\]

Produces GP sample paths that are twice differentiable.

Parameters:

lengthscale (Array) – Lengthscale controlling the smoothness.
priors (Dict[str, Callable], optional) – Priors for hyperparameters.
param_bounds (Dict)
exempt (Set[str])

lengthscale: Array¶

evaluate(x1, x2)¶

Evaluate the Matérn 5/2 kernel.

Parameters:

x1 (Array) – Input array of shape (n_samples_1, n_features).
x2 (Array) – Input array of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2).

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.SumKernel(kernel_1, kernel_2, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: KernelOperator

Represents a kernel that multiplies the outputs of two given kernels:

\[k(x, x'|\theta_1, \theta_2) = k_1(x, x'|\theta_1) + k_2(x, x'|\theta_2)\]

Parameters:

kernel_1 (Kernel)
kernel_2 (Kernel)
priors (Dict[str, Callable])
param_bounds (Dict)
exempt (Set[str])

kernel_1¶

The first kernel.

Type:: Kernel

kernel_2¶

The second kernel.

Type:: Kernel

kernel_1: Kernel¶

kernel_2: Kernel¶

evaluate(x1, x2)¶

Evaluate the kernel function.

Parameters:

x1 (Array) – First input data array.
x2 (Array) – Second input data array.

Returns:

Kernel evaluation between x1 and x2.

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.ProductKernel(kernel_1, kernel_2, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: KernelOperator

Represents a kernel that multiplies the outputs of two given kernels:

\[k(x, x'|\theta_1, \theta_2) = k_1(x, x'|\theta_1) + k_2(x, x'|\theta_2)\]

Parameters:

kernel_1 (Kernel)
kernel_2 (Kernel)
priors (Dict[str, Callable])
param_bounds (Dict)
exempt (Set[str])

kernel_1¶

The first kernel.

Type:: Kernel

kernel_2¶

The second kernel.

Type:: Kernel

kernel_1: Kernel¶

kernel_2: Kernel¶

evaluate(x1, x2)¶

Evaluate the kernel function.

Parameters:

x1 (Array) – First input data array.
x2 (Array) – Second input data array.

Returns:

Kernel evaluation between x1 and x2.

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.Scale(variance, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Scale Kernel.

The Scale kernel is a simple kernel that outputs a constant value across all pairs of points. It is often used in combination with other kernels to add a scaling factor.

The kernel is defined as:

\[k(\mathbf{x}_1, \mathbf{x}_2) = \sigma_f^2\]

Parameters:

variance (Array) – The variance (scaling factor) for the kernel.
priors (Dict[str, Callable], optional) – Priors for the hyperparameter (variance).
param_bounds (Dict)
exempt (Set[str])

variance: Array¶

evaluate(x1, x2)¶

Evaluate the Scale kernel between x1 and x2.

This kernel returns a constant value, which is the variance, for all pairs of points in x1 and x2.

Parameters:

x1 (Array) – Input array of shape (n_samples_1, n_features).
x2 (Array) – Input array of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2), where each element is the variance.

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.DiagNoise(noise, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Kernel that adds noise to the diagonal.

The kernel is defined as:

\[k(\mathbf{x}_i, \mathbf{x}_j) = \sigma_n^2 \delta_{ij}\]

Parameters:

variance (Array) – The variance (scaling factor) for the kernel.
priors (Dict[str, Callable], optional) – Priors for the hyperparameter (variance).
noise (Array)
param_bounds (Dict)
exempt (Set[str])

noise: Array¶

evaluate(x1, x2)¶

Evaluate the Scale kernel between x1 and x2.

This kernel returns a constant value, which is the variance, for all pairs of points in x1 and x2.

Parameters:

x1 (Array) – Input array of shape (n_samples_1, n_features).
x2 (Array) – Input array of shape (n_samples_2, n_features).

Returns:

Kernel matrix of shape (n_samples_1, n_samples_2), where each element is the variance.

Return type:

Array

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.SeparableKernel(kernel_1, kernel_2, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: KernelOperator

A separable kernel that factors into a product of a kernel over x and a kernel over w:

\[k((x, w), (x', w')) = k_x(x, x') \cdot k_w(w, w')\]

Parameters:

kernel_1 (Kernel)
kernel_2 (Kernel)
priors (Dict[str, Callable])
param_bounds (Dict)
exempt (Set[str])

kx¶: Kernel acting on the x inputs.

kw¶: Kernel acting on the w inputs.

evaluate(inputs1, inputs2)¶: inputs1: tuple (X, W) inputs2: tuple (X, W) X: shape (N_x, D_x) W: shape (N_w, D_w)

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

class ife_surrogate.gp.kernels.FixedFreqKernel(sigma, _C_w, *, priors, param_bounds=<factory>, exempt=<factory>)¶

Bases: Kernel

Kernel with a learnable spatial part Kx and a fixed frequency correlation matrix Cw. Kw = diag(alpha) @ Cw @ diag(alpha).

The frequency correlation is interpolated from a reference correlation matrix C_w defined on W_ref.

Parameters:

sigma (Array)
_C_w (Array)
priors (Dict[str, Callable])
param_bounds (Dict)
exempt (Set[str])

replace(**updates)¶: Returns a new object replacing the specified fields with new values.

sigma: Array¶

evaluate(W1, W2)¶: Interpolates Cw(W1, W2) from reference C_w.