This work is supported by GSoC, NumFOCUS, and PyMC team.

import math
import numpy as np
import pymc as pm
import arviz as az
import matplotlib.pyplot as plt
# set the seed
np.random.seed(1)
%matplotlib inline
%load_ext autoreload
%reload_ext autoreload
%autoreload 2

Set up training data

train_x = np.linspace(0, 1, 50)

train_y = np.stack([
    np.sin(train_x * (2 * math.pi)) + np.random.randn(len(train_x)) * 0.2,
    np.cos(train_x * (2 * math.pi)) + np.random.randn(len(train_x)) * 0.2,
    np.cos(train_x * (1 * math.pi)) + np.random.randn(len(train_x)) * 0.1,
], -1)

train_x.shape, train_y.shape

((50,), (50, 3))

fig, ax = plt.subplots(1,1, figsize=(12,5))
ax.scatter(train_x, train_y[:,0])
ax.scatter(train_x, train_y[:,1])
ax.scatter(train_x, train_y[:,2])
plt.legend(["sin", "cos"])

<matplotlib.legend.Legend at 0x7f8ac6e9bfa0>

x = train_x
xx = np.concatenate((x, x, x), axis=0)[:,None]
n = len(x)
idx2 = np.ones(n) + 1
idx = np.concatenate((np.zeros(n), np.ones(n), idx2))[:,None]
X = np.concatenate((xx, idx), axis=1)

y = np.concatenate((train_y[:,0], train_y[:,1], train_y[:,2]))
x.shape, X.shape, y.shape

((50,), (150, 2), (150,))

LCM model in PyMC

X.shape, y.shape

((150, 2), (150,))

with pm.Model() as model:
    ell = pm.Gamma("ell", alpha=2, beta=0.5)
    eta = pm.Gamma("eta", alpha=2, beta=0.5)
    cov = eta**2 * pm.gp.cov.ExpQuad(2, ls=ell, active_dims=[0])
    
    ell2 = pm.Gamma("ell2", alpha=2, beta=0.5)
    eta2 = pm.Gamma("eta2", alpha=2, beta=0.5)
    cov2 = eta**2 * pm.gp.cov.Matern32(2, ls=ell, active_dims=[0])
    
    W = pm.Normal("W", mu=0, sigma=3, shape=(3,2), initval=np.random.randn(3,2))
    kappa = pm.Gamma("kappa", alpha=1.5, beta=1, shape=3)
    coreg = pm.gp.cov.Coregion(input_dim=2, active_dims=[1], kappa=kappa, W=W)
    
    W2 = pm.Normal("W2", mu=0, sigma=3, shape=(3,2), initval=np.random.randn(3,2))
    kappa2 = pm.Gamma("kappa2", alpha=1.5, beta=1, shape=3)
    coreg2 = pm.gp.cov.Coregion(input_dim=2, active_dims=[1], kappa=kappa2, W=W2)
    
    cov_func1 = coreg * cov #pm.gp.cov.Prod([coreg, cov])
    cov_func2 = coreg2 * cov2 #pm.gp.cov.Prod([coreg2, cov2])
    cov_func = cov_func1 + cov_func2 #pm.gp.cov.Add([cov_func1, cov_func2])
    
    sigma = pm.HalfNormal("sigma", sigma=3)
    gp = pm.gp.Marginal(cov_func=cov_func)
    y_ = gp.marginal_likelihood("f", X, y, noise=sigma)

%%time
with model:
    gp_trace = pm.sample(500, chains=1)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Sequential sampling (1 chains in 1 job)
NUTS: [ell, eta, ell2, eta2, W, kappa, W2, kappa2, sigma]

Sampling 1 chain for 1_000 tune and 500 draw iterations (1_000 + 500 draws total) took 246 seconds.
There was 1 divergence after tuning. Increase `target_accept` or reparameterize.

CPU times: user 11min 13s, sys: 21min 26s, total: 32min 39s
Wall time: 4min 16s

x_new = np.linspace(-0.5, 1.5, 200)[:, None]
xx_new = np.concatenate((x_new, x_new, x_new), axis=0)
idx2 = np.ones(200) + 1
idx2 = np.concatenate((np.zeros(200), np.ones(200), idx2))[:, None]
X_new = np.concatenate((xx_new, idx2), axis=1)

X_new.shape

(600, 2)

with model:
    preds = gp.conditional("preds", X_new)
    gp_samples = pm.sample_posterior_predictive(gp_trace, var_names=['preds'], random_seed=42)

from pymc.gp.util import plot_gp_dist
fig = plt.figure(figsize=(12,5))
ax = fig.gca()

f_pred = gp_samples.posterior_predictive["preds"].sel(chain=0)
plot_gp_dist(ax, f_pred[:,:200], X_new[:200,0], palette="Blues", fill_alpha=0.5, samples_alpha=0.1)
ax.plot(x, train_y[:,0], 'ok', ms=3, alpha=0.5, label="Data 1");

from pymc.gp.util import plot_gp_dist
fig = plt.figure(figsize=(12,5))
ax = fig.gca()

plot_gp_dist(ax, f_pred[:,200:400], X_new[200:400,0], palette="Blues", fill_alpha=0.9, samples_alpha=0.1)
ax.plot(x, train_y[:,1], 'ok', ms=3, alpha=0.5, label="Data 2");
ax.set_ylim([-4,4])

(-4.0, 4.0)

from pymc.gp.util import plot_gp_dist
fig = plt.figure(figsize=(12,5))
ax = fig.gca()

plot_gp_dist(ax, f_pred[:,400:], X_new[400:,0], palette="Blues", fill_alpha=0.9, samples_alpha=0.1)
ax.plot(x, train_y[:,2], 'ok', ms=3, alpha=0.5, label="Data 2");
ax.set_ylim([-4,4])

(-4.0, 4.0)

az.summary(gp_trace)

arviz - WARNING - Shape validation failed: input_shape: (1, 500), minimum_shape: (chains=2, draws=4)

az.plot_trace(gp_trace);

%load_ext watermark
%watermark -n -u -v -iv -w

Last updated: Sun Sep 04 2022

Python implementation: CPython
Python version       : 3.9.12
IPython version      : 8.3.0

numpy     : 1.22.4
matplotlib: 3.5.2
arviz     : 0.12.1
pymc      : 4.1.5

Watermark: 2.3.0

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
W[0, 0]	0.073	2.862	-5.200	5.049	0.133	0.131	471.0	395.0	NaN
W[0, 1]	-0.022	2.673	-5.814	4.383	0.121	0.114	488.0	395.0	NaN
W[1, 0]	-0.170	2.883	-5.834	4.848	0.146	0.140	388.0	361.0	NaN
W[1, 1]	0.091	3.000	-5.205	5.548	0.138	0.139	477.0	409.0	NaN
W[2, 0]	-0.346	2.614	-5.511	4.344	0.149	0.116	307.0	311.0	NaN
W[2, 1]	0.011	2.646	-4.982	4.636	0.133	0.094	413.0	439.0	NaN
W2[0, 0]	0.094	3.283	-5.072	6.234	0.215	0.152	239.0	374.0	NaN
W2[0, 1]	0.174	3.597	-6.100	6.183	0.244	0.173	226.0	369.0	NaN
W2[1, 0]	0.096	3.332	-5.530	5.907	0.242	0.171	192.0	313.0	NaN
W2[1, 1]	-0.032	3.101	-4.922	6.101	0.220	0.156	199.0	408.0	NaN
W2[2, 0]	0.004	0.914	-1.628	1.598	0.056	0.040	261.0	315.0	NaN
W2[2, 1]	0.034	0.992	-1.661	2.019	0.057	0.043	301.0	280.0	NaN
ell	0.769	0.536	0.224	1.720	0.042	0.030	218.0	221.0	NaN
eta	0.500	0.527	0.108	1.240	0.038	0.027	224.0	248.0	NaN
ell2	3.982	2.831	0.173	8.617	0.125	0.089	375.0	305.0	NaN
eta2	3.867	2.728	0.032	8.138	0.133	0.094	247.0	95.0	NaN
kappa[0]	1.476	1.231	0.054	3.753	0.047	0.036	562.0	295.0	NaN
kappa[1]	1.476	1.254	0.001	3.782	0.050	0.036	316.0	189.0	NaN
kappa[2]	1.499	1.149	0.016	3.520	0.047	0.033	404.0	260.0	NaN
kappa2[0]	1.655	1.357	0.008	4.084	0.058	0.043	292.0	206.0	NaN
kappa2[1]	1.636	1.244	0.041	3.758	0.048	0.034	480.0	303.0	NaN
kappa2[2]	1.091	0.988	0.023	2.943	0.041	0.030	343.0	244.0	NaN
sigma	0.153	0.010	0.133	0.169	0.000	0.000	602.0	333.0	NaN