Oak

Orthogonal Additive Kernel (OAK).

Reference

Lu, X., Boukouvalas, A., & Hensman, J. (2022). Additive Gaussian Processes Revisited. ICML.

OrthogonalAdditiveKernel

OrthogonalAdditiveKernel(
    base_kernels: list[AbstractKernel],
    max_order: Union[int, None] = None,
    order_variances: Union[
        Float[Array, " D_tilde_plus_1"], None
    ] = None,
    fix_base_variance: bool = True,
    compute_engine: AbstractKernelComputation = DenseKernelComputation(),
)

Bases: AbstractKernel

Orthogonal Additive Kernel (OAK).

Wraps D one-dimensional SE base kernels with an orthogonality constraint (under standard normal input density) and combines them via Newton-Girard into an additive kernel with configurable maximum interaction order.

The kernel decomposes as

K = sum_{l=0}^{D_tilde} sigma^2_l * E_l

where E_l is the l-th elementary symmetric polynomial of the D constrained base kernel evaluations, and sigma^2_l are learnable order variances.

Reference

Lu, X., Boukouvalas, A., & Hensman, J. (2022). Additive Gaussian Processes Revisited. ICML.

Parameters:

• base_kernels (list[AbstractKernel]) –

  List of D one-dimensional base kernels (typically RBF with active_dims=[i] for each dimension i). Each must have lengthscale and variance attributes.

• max_order (Union[int, None], default: None ) –

  Maximum interaction order (D_tilde). Defaults to D. Must be <= D.

• order_variances (Union[Float[Array, ' D_tilde_plus_1'], None], default: None ) –

  Initial order variances of shape (max_order + 1,). Entry 0 is the offset variance, entry d is the d-th order interaction variance. Defaults to ones.

• fix_base_variance (bool, default: True ) –

  If True (default), pin every base-kernel variance to 1 so that order_variances alone control per-order scaling. This avoids over-parameterisation and matches the reference (Lu et al. 2022, S3.2).

• compute_engine (AbstractKernelComputation, default: DenseKernelComputation() ) –

  Kernel computation engine. Defaults to DenseKernelComputation.

__call__

__call__(
    x: Float[Array, " D"], y: Float[Array, " D"]
) -> ScalarFloat

Evaluate the OAK kernel on a pair of inputs.

Computes constrained kernel values per dimension via vmap, then combines via Newton-Girard recursion weighted by order variances.

Parameters:

• x (Float[Array, ' D']) –

  Left input of shape (D,).

• y (Float[Array, ' D']) –

  Right input of shape (D,).

Returns:

• ScalarFloat –

  Scalar kernel value.

cross_covariance

cross_covariance(
    x: Num[Array, "N D"], y: Num[Array, "M D"]
) -> Float[Array, "N M"]

Compute the cross-covariance matrix of the kernel.

Parameters:

• x (Num[Array, 'N D']) –

  the first input matrix of shape (N, D).

• y (Num[Array, 'M D']) –

  the second input matrix of shape (M, D).

Returns:

• Float[Array, 'N M'] –

  The cross-covariance matrix of the kernel of shape (N, M).

gram

gram(x: Num[Array, 'N D']) -> LinearOperator

Compute the gram matrix of the kernel.

Parameters:

• x (Num[Array, 'N D']) –

  the input matrix of shape (N, D).

Returns:

• LinearOperator –

  The gram matrix of the kernel of shape (N, N).

diagonal

diagonal(x: Num[Array, 'N D']) -> LinearOperator

Compute the diagonal of the gram matrix of the kernel.

Parameters:

• x (Num[Array, 'N D']) –

  the input matrix of shape (N, D).

Returns:

• LinearOperator –

  The diagonal of the gram matrix of the kernel of shape (N,).

slice_input

slice_input(
    x: Float[Array, "... D"],
) -> Float[Array, "... Q"]

Slice out the relevant columns of the input matrix.

Select the relevant columns of the supplied matrix to be used within the kernel's evaluation.

Parameters:

• x (Float[Array, '... D']) –

  the matrix or vector that is to be sliced.

Returns:

• Float[Array, '... Q'] –

  The sliced form of the input matrix.

__add__

__add__(
    other: Union[AbstractKernel, ScalarFloat],
) -> AbstractKernel

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns:

• AbstractKernel ( AbstractKernel ) –

  A new kernel that is the sum of the two kernels.

__mul__

__mul__(
    other: Union[AbstractKernel, ScalarFloat],
) -> AbstractKernel

Multiply two kernels together.

Parameters:

• other (AbstractKernel) –

  The kernel to be multiplied with the current kernel.

Returns:

• AbstractKernel ( AbstractKernel ) –

  A new kernel that is the product of the two kernels.