A Julia package for computing 2D GPs with Kronecker product kernels. Based on luas from Fortune et al. 2024. This package uses the KernelFunctions.jl and AbstractGPs.jl interfaces, which works with inference tools such as Turing.jl. Kernel matrix evaluations rely on Kronecker.jl for fast computation and low-memory overhead.
using AbstractGPs, KernelFunction, KroneckerProductKernels
using LinearAlgebra, Random; Random.seed!(1234)
using Kronecker: ⊗ # \otimes<tab>
# Create a KernelKroneckerProduct using standard KernelFunctions kernels
k1 = 0.1 * Matern52Kernel()
k2 = SqExponentialKernel() ∘ ScaleTransform(0.4) # \circ<tab> -> ∘
kernel = KernelKroneckerProduct(k1, k2)
# And get an AbstractGP
f = GP(kernel)
# Inputs must be structured in a particular way
# Need a set of 2nd dimension values for every 1st dimension value
# Inputs should be sorted by the first dimension
X1 = collect(0.0:1:100)
X2 = collect(0.0:1:5)
X = vcat([[x1 x2] for x1 in X1 for x2 in X2]...)
# We can now get a FiniteGP
# Either with a uniform, diagonal covariance
fx = f(RowVecs(X), 0.1)
# Or with a non-uniform, diagonal covariance
# This must be provided as a KroneckerProduct
C1 = Diagonal(rand(length(X1)))
C2 = Diagonal(rand(length(X2)))
C = C1 ⊗ C2
fx = f(RowVecs(X), C)
# Now, we can draw from the GP
Y = rand(fx)
# And evaluate the logpdf
logpdf(fx, Y)