Short example of a space/time Kronecker Gaussian process in PyMC3

spacetimekron.py

import numpy as np
import pymc3 as pm

N, T = 20, 10

min_lat, max_lat   = 35, 40
min_long, max_long = 30, 35 

lat_interval  = max_lat - min_lat
long_interval = max_long - min_long

spatial_coordinates  = np.random.uniform(size=[N,2]) * np.asarray([[long_interval,lat_interval]])
spatial_coordinates += np.asarray([[min_long, min_lat]])
temporal_coordinates = np.linspace(0,5,T)

y = np.random.randn(N*T)

kron_Xs = [spatial_coordinates[:,[0,1]], temporal_coordinates[:, np.newaxis]]

with pm.Model() as model:
    
    # Parameters governing the correlation distances and magnitudes
    # for different components of this model
    ls        = pm.Uniform('ls', lower=0.2, upper=4.0)
    ls_period = pm.Uniform('ls_period', lower=0.1, upper=4.0)
    ls_time   = pm.Uniform('ls_time', lower=0.2, upper=4.0)
    scales    = pm.InverseGamma('gp_variances', alpha=1.0, beta=1.0, shape=4)**0.5
    
    # Spatial covariance 
    matern_space   = pm.gp.cov.Matern52(2, ls=ls, active_dims=[0,1]) * scales[0]

    # Temporal covariance expressed as the sum of a local smoothing (Matern)
    # and a sinusoidal / repeating function (periodic)
    periodic    = pm.gp.cov.Cosine(1, ls=ls_period, active_dims=[0]) * scales[1]
    matern_time = pm.gp.cov.Matern52(1, ls=ls_time, active_dims=[0]) * scales[2]

    # White noise term 
    sigma    = pm.HalfNormal('sigma')
     
    gp = pm.gp.MarginalKron(cov_funcs=[matern_space, periodic + matern_time])
    ll = gp.marginal_likelihood('ll', kron_Xs, y, sigma=sigma, is_observed=True)
    trace = pm.sample(chains=1, cores=1, tune=300, draws = 300)