Skip to content

Instantly share code, notes, and snippets.

@ckrapu
Created February 21, 2021 18:06
Show Gist options
  • Save ckrapu/1d7e2e76ce3f228b19029e4fecc9acbf to your computer and use it in GitHub Desktop.
Save ckrapu/1d7e2e76ce3f228b19029e4fecc9acbf to your computer and use it in GitHub Desktop.
Short example of a space/time Kronecker Gaussian process in PyMC3
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment