Skip to content

Kernel

MaternKernel

gptempest.MaternKernel

Bases: Module

Matérn covariance kernel for Gaussian Process regression.

Computes the Matérn kernel between sets of time points. The smoothness parameter ν controls the differentiability of the GP sample paths: ν=0.5 gives the Ornstein–Uhlenbeck process, ν=1.5 once-differentiable paths, and ν=2.5 twice-differentiable paths.

Parameters:

Name Type Description Default
scale float

Time-scale parameter. Larger values mean slower-varying GPs. Typically set to the number of frames in the trajectory.

required
nu float

Matérn smoothness parameter. Must be one of 0.5, 1.5, or 2.5.

required
dtype dtype

PyTorch dtype, e.g. torch.float64.

required

Raises:

Type Description
RuntimeError

If nu is not one of the supported values.

Example
kernel = MaternKernel(nu=1.5, scale=1000.0, dtype=torch.float64)
K = kernel.kernel_mat(t1, t2)  # shape (N, M)
Source code in src/gptempest/model.py
class MaternKernel(nn.Module):
    """Matérn covariance kernel for Gaussian Process regression.

    Computes the Matérn kernel between sets of time points. The smoothness
    parameter ν controls the differentiability of the GP sample paths:
    ν=0.5 gives the Ornstein–Uhlenbeck process, ν=1.5 once-differentiable
    paths, and ν=2.5 twice-differentiable paths.

    Args:
        scale: Time-scale parameter. Larger values mean slower-varying GPs.
            Typically set to the number of frames in the trajectory.
        nu: Matérn smoothness parameter. Must be one of ``0.5``, ``1.5``,
            or ``2.5``.
        dtype: PyTorch dtype, e.g. ``torch.float64``.

    Raises:
        RuntimeError: If ``nu`` is not one of the supported values.

    Example:
        ```python
        kernel = MaternKernel(nu=1.5, scale=1000.0, dtype=torch.float64)
        K = kernel.kernel_mat(t1, t2)  # shape (N, M)
        ```
    """
    def __init__(self, scale: float, nu: float, dtype: torch.dtype):
        super(MaternKernel, self).__init__()
        self.scale = torch.tensor([scale], dtype=dtype)
        self.device = self.scale.device
        self.dtype = dtype
        self.nu = nu
        if nu not in {0.5, 1.5, 2.5}:
            raise RuntimeError('nu expected to be 0.5, 1.5, or 2.5')

    def to(self, device):
        """Move the kernel and its parameters to the specified device."""
        self.device = device
        self.scale = self.scale.to(device)
        return self

    def _compute_kernel(self, distance):
        """Evaluate the Matérn kernel for pre-scaled pairwise distances.

        Args:
            distance: Tensor of non-negative scaled distances.

        Returns:
            Tensor of kernel values with the same shape as ``distance``.
        """
        exp_component = torch.exp(-torch.sqrt(torch.tensor(
            self.nu * 2, dtype=self.dtype, device=self.device
        )) * distance)
        if self.nu == 0.5:
            prefac = 1
        elif self.nu == 1.5:
            prefac = (torch.sqrt(
                torch.tensor(3.0, dtype=self.dtype, device=self.device)
            ) * distance).add(1)
        elif self.nu == 2.5:
            prefac = (torch.sqrt(
                torch.tensor(5.0, dtype=self.dtype, device=self.device)
            ) * distance).add(1).add(5.0 / 3.0 * distance**2)
        return prefac * exp_component

    def _scale_times(self, t1, t2):
        """Center and scale two time arrays by the kernel time-scale.

        Subtracts the mean of ``t1`` from both arrays and divides by
        ``self.scale``, ensuring the GP is not sensitive to the absolute
        position of the trajectory in time.

        Args:
            t1: Time array of shape ``(N, 1)``.
            t2: Time array of shape ``(M, 1)``.

        Returns:
            Tuple ``(t1_scaled, t2_scaled)`` of the same shapes.
        """
        mean = t1.mean(dim=-2, keepdim=True)
        t1s = t1.sub(mean).div_(self.scale)
        t2s = t2.sub(mean).div_(self.scale)
        return t1s, t2s

    def kernel_mat(self, t1, t2):
        """Compute the full Matérn kernel matrix between two time arrays.

        Args:
            t1: Time array of shape ``(N, 1)``.
            t2: Time array of shape ``(M, 1)``.

        Returns:
            Kernel matrix of shape ``(N, M)``.
        """
        t1_s, t2_s = self._scale_times(t1, t2)
        distance = torch.cdist(t1_s, t2_s)
        return self._compute_kernel(distance)

    def kernel_diag(self, t1, t2):
        """Compute element-wise Matérn kernel values between paired time points.

        Equivalent to the diagonal of ``kernel_mat(t1, t2)`` but more
        efficient as it avoids computing the full pairwise distance matrix.

        Args:
            t1: Time array of shape ``(N, 1)``.
            t2: Time array of shape ``(N, 1)``.

        Returns:
            Kernel values of shape ``(N,)``.
        """
        t1_s, t2_s = self._scale_times(t1, t2)
        distance = ((t1_s - t2_s)**2).sum(dim=1).sqrt()
        return self._compute_kernel(distance).squeeze()

to(device)

Move the kernel and its parameters to the specified device.

Source code in src/gptempest/model.py
def to(self, device):
    """Move the kernel and its parameters to the specified device."""
    self.device = device
    self.scale = self.scale.to(device)
    return self

kernel_mat(t1, t2)

Compute the full Matérn kernel matrix between two time arrays.

Parameters:

Name Type Description Default
t1

Time array of shape (N, 1).

required
t2

Time array of shape (M, 1).

required

Returns:

Type Description

Kernel matrix of shape (N, M).

Source code in src/gptempest/model.py
def kernel_mat(self, t1, t2):
    """Compute the full Matérn kernel matrix between two time arrays.

    Args:
        t1: Time array of shape ``(N, 1)``.
        t2: Time array of shape ``(M, 1)``.

    Returns:
        Kernel matrix of shape ``(N, M)``.
    """
    t1_s, t2_s = self._scale_times(t1, t2)
    distance = torch.cdist(t1_s, t2_s)
    return self._compute_kernel(distance)

kernel_diag(t1, t2)

Compute element-wise Matérn kernel values between paired time points.

Equivalent to the diagonal of kernel_mat(t1, t2) but more efficient as it avoids computing the full pairwise distance matrix.

Parameters:

Name Type Description Default
t1

Time array of shape (N, 1).

required
t2

Time array of shape (N, 1).

required

Returns:

Type Description

Kernel values of shape (N,).

Source code in src/gptempest/model.py
def kernel_diag(self, t1, t2):
    """Compute element-wise Matérn kernel values between paired time points.

    Equivalent to the diagonal of ``kernel_mat(t1, t2)`` but more
    efficient as it avoids computing the full pairwise distance matrix.

    Args:
        t1: Time array of shape ``(N, 1)``.
        t2: Time array of shape ``(N, 1)``.

    Returns:
        Kernel values of shape ``(N,)``.
    """
    t1_s, t2_s = self._scale_times(t1, t2)
    distance = ((t1_s - t2_s)**2).sum(dim=1).sqrt()
    return self._compute_kernel(distance).squeeze()