Implementing a GP for spatio-temporal data with Kronecker structured covariance

[fbartolic]

Hello! I have a spatio-temporal dataset where each snapshot in time is evaluated on a 13 X 13 regular spatial grid and the number of (irregularly spaced) time observations is very large (>100k). I want to assume Kronecker structure for the covariance of the form K = k_space(x, x') kron k_time(t, t') where k_time is non-stationary. For now, I'm trying to get a simpler model to work using GridInterpolationKernel for the time covariance but I'm running into shape issues and I'm not sure how to proceed.

The code looks like this:

# Spatial grid
train_x = gpytorch.utils.grid.create_data_from_grid([train_t, xgrid, ygrid])

# Standardize the data
train_y = (train_y - torch.mean(train_y))/torch.std(train_y)

class GPRegression(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):   
        super(GPRegression, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ZeroMean()

        # Time covariance
        self.covar_time = gpytorch.kernels.GridInterpolationKernel(
            gpytorch.kernels.ScaleKernel(
                gpytorch.kernels.RBFKernel(active_dims=[0])
            ), 
            grid_size=len(train_t),
            num_dims=1
        )
        
        # Spatial covariance
        self.covar_space = gpytorch.kernels.GridKernel(
            gpytorch.kernels.ScaleKernel(RBFKernel(active_dims=[1, 2])),
             grid=torch.stack([xgrid, ygrid]).T
         )

        self.likelihood = likelihood

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = gpytorch.lazy.KroneckerProductLazyTensor(self.covar_time(x),  self.covar_space(x))
        
        return gpytorch.distributions.MultivariateNormal(mean_x,  covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = GPRegression(train_x, train_y, likelihood)

where train_x has shape [Ntotal, 3] and train_y has shape [Ntotal] with Ntotal = Ntime*Nspace. This seems closely related to #319 but I still wasn't able to get it to work. Any help would be appreciated!

[gpleiss]

covar_x = gpytorch.lazy.KroneckerProductLazyTensor(self.covar_time(x), self.covar_space(x)) is going to cause trouble. Three suggestions:

1. covar_x = self.covar_time(x) * self.covar_space(x). This won't be the most efficient, but it will still do some structure exploitation. It will implement the technique used in the SKIP paper: https://arxiv.org/abs/1802.08903
2. Use the GridInterpolationKernel for space and time dimensions. This will use Kronecker structure under the hood, and - for all intents and purposes - will be just as accurate as if you weren't using interpolation for the time dimension.
3. Implementing your own kernel that does interpolation for spatial dimensions and non-interpolation for the time dimension. This last option probably isn't worth it, in my opinion.