Source code for gpjax.kernels.multioutput.icm

from gpjax.kernels.base import AbstractKernel
from gpjax.kernels.multioutput.base import MultiOutputKernel
from gpjax.kernels.multioutput.computation import MultiOutputKernelComputation
from gpjax.parameters import CoregionalizationMatrix




class ICMKernel(MultiOutputKernel):
    """Intrinsic Coregionalization Model kernel.

    Combines a single shared input kernel with a coregionalization matrix
    to produce a Kronecker-structured covariance: K = B (x) K_input.

    Args:
        base_kernel: The shared input-space kernel.
        coregionalization_matrix: The output-space coregionalization.
    """

    base_kernel: AbstractKernel
    coregionalization_matrix: CoregionalizationMatrix

    def __init__(
        self,
        base_kernel: AbstractKernel,
        coregionalization_matrix: CoregionalizationMatrix,
    ):
        self.base_kernel = base_kernel
        self.coregionalization_matrix = coregionalization_matrix
        super().__init__(compute_engine=MultiOutputKernelComputation())

    @property
    def num_outputs(self) -> int:
        return self.coregionalization_matrix.num_outputs

    @property
    def num_latent_gps(self) -> int:
        return 1

    @property
    def latent_kernels(self) -> tuple[AbstractKernel, ...]:
        return (self.base_kernel,)

    @property
    def components(self):
        return ((self.coregionalization_matrix, self.base_kernel),)