Here’s maybe a better example with a custom noise kernel:
import numpy as np
import pymc3 as pm
import theano.tensor as tt
class Heteroskedastic(pm.gp.cov.Covariance):
## adapted from pm.gp.cov.WhiteNoise
def __init__(self, sigmas):
super().__init__(1, None)
self.sigmas = tt.as_tensor_variable(sigmas)
def diag(self, X):
return self.sigmas
def full(self, X, Xs=None):
if Xs is None:
return tt.diag(self.diag(X))
else:
return tt.alloc(0.0, X.shape[0], Xs.shape[0])
## underlying process
x = np.linspace(0, 10, 100)
f = np.sin(2 * np.pi * x * 0.1)
## heteroskedastic noise
sigma = 0.15 * x + 0.25
noise = sigma * np.random.randn(100)
y = f + noise
plt.plot(x, f)
plt.plot(x, y, '.k');
## fit model
with pm.Model() as model:
eta = pm.HalfNormal('eta', sd=5)
ell = 3
cov = eta**2 * pm.gp.cov.ExpQuad(1, ls=ell)
gp = pm.gp.Marginal(cov_func=cov)
a = pm.Normal('a', mu=0.0, sd=5.0)
b = pm.HalfNormal('b', sd=5.0)
sigmas = pm.Deterministic('sigmas', a*x + b)
noise = Heteroskedastic(sigmas)
gp.marginal_likelihood("y", X=x[:, None], y=y, noise=noise)
tr = pm.sample(cores=1, chains=1)
## show fit
with model:
f = gp.conditional('f', Xnew=x[:, None])
ppc = pm.sample_posterior_predictive(trace=tr, var_names=['f', 'y'])
ub = np.percentile(ppc['f'], 95, axis=0)
lb = np.percentile(ppc['f'], 5, axis=0)
plt.plot(x, y, '.k')
plt.fill_between(x, lb, ub, color='slateblue', alpha=0.5)