Passing a static matrix with covariance kernel for pm.gp.LatentKron

Hi,

I have a model where a set of latent RVs are defined using a GP. The index of these RVs is (d, h), where d denotes day of week ({1, 2, …, 7}) and h denotes hours ({0, 1, …, 23}). I would like to specify the the full covariance matrix as the Kronecker product of two covariance matrices, where matrix 1) is some arbitrary covariance matrix with dimensions (dxd), and matrix 2) is specifed by a kernel function that operates on the index h.

It appears that pm.gp.LatentKron can only accept kernel functions in cov_funcs. Is it possible to pass in a covariance matrix that is defined without using the kernel functions?

To write what I’d like to be able to do specifically:

# Define kernel function for hourly index

cov_hourly = pm.gp.cov.Periodic(input_dim=1, period=24, ls=1/2)

# Define covariance matrix for daily index with an LKJ prior

mu_omega_chol, mu_omega_corr, mu_omega_stds = pm.LKJCholeskyCov(
   "mu_omega_chol", n=7, eta=2.0, sd_dist=pm.Exponential.dist(1.0), compute_corr=True
)

cov_daily = pm.Deterministic("mu_omega_cov", mu_omega_chol.dot(mu_omega_chol.T))

# Combine to pm.gp.LatentKron (Illegal to do this)

gp = pm.gp.LatentKron(cov_funcs=[cov_daily, cov_hourly])

Please let me know if something needs further clarification.

You may need to make a thin covariance function shim to go around cov_daily. For instance, here’s the code for pm.gp.cov.Constant which does this for a scalar value.

import aesara.tensor as at

class Constant(pm.gp.cov.Covariance):
    r"""
    Constant valued covariance function.

    .. math::

       k(x, x') = c
    """

    def __init__(self, c):
        super().__init__(1, None)
        self.c = c

    def diag(self, X):
        return at.alloc(self.c, X.shape[0])

    def full(self, X, Xs=None):
        if Xs is None:
            return at.alloc(self.c, X.shape[0], X.shape[0])
        else:
            return at.alloc(self.c, X.shape[0], Xs.shape[0])

The extension of this for a covariance matrix should be even simpler. Here’s a little bit more info on how to do this: Mean and Covariance Functions — PyMC3 3.11.5 documentation

2 Likes

Thanks @bwengals. Interesting idea, I will try it out.