Porting a Pymc3 model to pyro (latent Gaussian Process)

Misc.
1

I am a newbie in pyro, so please bear with me in case I miss anything obvious.

I have a generative model for a group of time series that is basically a single latent Gaussian Process summed with a number of piecewise linear functions equal to the number of time series. I want to know if this is possible to implement in pyro and if I should expect performance gains.

Here you can find an example of the generative process and plot the series:

import pymc3 as pm
import numpy as np
import matplotlib.pyplot as plt
import theano
import theano.tensor as tt

class OutPiecewiseLinearChangepoints():
    # Piecewise linear function applied outside of the GPs and added in the likelihood parameter
    def __init__(self, 
                 k, 
                 m,
                 b,
                 changepoints,
                 n_series = 4):
        self.k = k
        self.m = m
        self.b = b
        self.changepoints = changepoints
        self.n_series = n_series

    def create_changepoints(self, X, changepoints):
        return (0.5 * (1.0 + tt.sgn(tt.tile(X.reshape((-1,1)), (1,len(changepoints))) - changepoints)))

    def build(self, X):

        X = theano.shared(X)

        A = self.create_changepoints(X, self.changepoints)
        
        growth = (self.k.reshape((1, -1)) + tt.dot(A, self.b))*tt.tile(X, (1, self.n_series))
        offset = (self.m.reshape((1,-1)) + tt.dot(A, (-self.changepoints.reshape((-1,1)) * self.b)))
        piecewise = growth + offset

        return piecewise

    piece = OutPiecewiseLinearChangepoints(k=np.array([0.1]), 
                                   m=np.array([0.1]), 
                                   b=np.random.normal(0.1, 0.4, size=(4,4)), 
                                   changepoints = changepoints_t,
                                   ).build(np.arange(10).reshape((-1,1))).eval()

    lengthscale = 0.2
    eta = 2.0
    cov = eta ** 2 * pm.gp.cov.ExpQuad(1, lengthscale)

    X = np.arange(10)[:, None]
    K = cov(X).eval()

    gpl = pm.MvNormal.dist(mu=np.zeros(K.shape[0]), cov=K).random(size=10).T

    x_train = np.random.normal(pm.MvNormal.dist(mu=np.zeros(K.shape[0]), cov=K).random(size=4).T + piece, 0.1)
    plt.figure(figsize=(14, 4))
    plt.plot(X, x_train)
    plt.title("time series generated")
    plt.ylabel("y")
    plt.xlabel("X");

With pymc3 I can build the model as:

with pm.Model() as model:
    period = pm.Gamma('period', 40, 10)

    l_t = pm.InverseGamma('l_t', 4, 40)
    l_p = pm.HalfNormal('l_p', 0.5)
    η_trend = pm.HalfNormal('η_trend',1)
    σ  = pm.HalfNormal("σ",  sigma=0.1)

    mu_func = pm.gp.mean.Zero()

    # cov function for the GP 
    cov = (η_trend**2 * pm.gp.cov.ExpQuad(input_dim=1, ls=l_t)
            + pm.gp.cov.WhiteNoise(σ))

    gp = pm.gp.Latent(mean_func=mu_func, cov_func=cov)
    f = gp.prior('f', X=X, reparameterize=True)

    k = pm.Normal("k", mu=0, sigma=0.01, shape = s)
    m = pm.Normal("m", mu=0, sigma=0.01, shape = s)
    delta = pm.Normal("delta", mu=0, sigma=0.01, shape=(len(changepoints_t),s))
    
    piece = OutPiecewiseLinearChangepoints(k = k,
                                          m = m,
                                          b = delta,
                                          changepoints = changepoints_t).build(X)

    # Likelihood
    y_pred = pm.Normal('y_pred', mu=f.reshape((-1,1)) + piece, sigma = σ,observed=x_train)

    # Fitting using ADVI
    trace_vi = pm.fit(100000,
                       method='advi',
                       obj_n_mc=1,
                       obj_optimizer=pm.adamax(),
                       callbacks=[pm.callbacks.CheckParametersConvergence(every=50, diff='absolute',tolerance=1e-3)],
                       progressbar = True)
    
    # Sampling
    trace_vi_samples = trace_vi.sample()
    pred_samples_fit = pm.sample_posterior_predictive(trace_vi_samples,
                                                           vars=[y_pred],
                                                           samples=n_samples,
                                                      progressbar = True)
    
    # Conditional of the GP on new points
    f_n = gp.conditional('f_n', Xnew=X_new)

    piece = OutPiecewiseLinearChangepoints(k = k,
                                          m = m,
                                          b = delta,
                                          changepoints = changepoints_t).build(X_new)

    
    y_pred_new = pm.Normal("y_pred_new", 
                            mu=f_n.reshape((-1,1)) + piece, 
                            shape=(X_new.shape[0], s))
    
    # Sampling of predictions
    pred_samples_predict = pm.sample_posterior_predictive(trace_vi_samples, 
                                              vars=[y_pred_new], 
                                              samples=n_samples,
                                              progressbar = True)

Once again, Is this possible to implement in pyro? Should I expect performance gains with such a model in pyro?

Thanks a lot!
luis

2

The answers to these questions are, as usual, “Yes, but…” and “it depends.” If I were you I’d want a clear motivation other than performance for rewriting the code in Pyro, like access to a Pyro-specific feature or compatibility with the larger PyTorch ecosystem. You might start by reviewing some of our GP examples or GPyTorch, which is compatible with Pyro.

1 Like
3

Thanks @eb8680_2, I checked it and this is promising (for future reference)!

1 Like