Heteroskedastic Marginal Likelihood GP

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)

image

4 Likes