Combining Gaussian process with mixed effects

[sadatnfs]

Hello,

I was wondering if it's possible to maybe implement custom mean functions in a GPR model? The use case I have is to fit a model of the following type:

Y_{c,t} = alpha + alpha_c + beta * X_{c,t} + (Time GP_t)
alpha_c ~ Normal(0., sigma_a)
GP_t ~ Matern52(., .)

Therefore, the mean function I'd like to feed into the GP would be alpha + alpha_c + beta*X, where alpha_c is a random effect on the c dimension, and the GP is on the t dimension.

The reason why I'd like to fit a model like this is to be able to model a mixed effects regression with GP errors. Using Edward2 from TFP, I can do something like tf.gather(parameter_vector, vector_of_indices_indicating_different_c), which will do like a matrix multiplication of the index and the associated parameter we're fitting on the random effect. I believe I can follow through on the various mixed effect models examples you guys have, so the last piece that I'd need to implement there would be to add a GP in the model.

Two things I can think of two implement this would be:

1. Building a GPR where the mean function will have all the traditional random effects and multiple fixed effects as needed.
2. Using the linear mixed effects model example, where I build in a GP component in the model.

I think the second would be preferable just because it'd give us more flexibility overall I think. Ideally, if there's a way to add a Gaussian process term in the model here, then that'd be the best: https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Linear_Mixed_Effects_Models.ipynb

Any tips would be great! Thanks!

[csuter]

Hi @sadatnfs, i haven't thought too deeply about this model in particular, but have you seen that there is a mean_fn argument on the GP classes? You can pass one in and it will be evaluated at the given index_points (and observation_index_points in the GPRModel). Does this do the trick?

[sadatnfs]

Thanks @csuter. I think that actually might work in the optimization process! What I don't know yet is what would be the best way to do posterior sampling of the parameters. As far as I know, the GaussianProcessRegressionModel module can only sample from the posterior parameters of the GP, and not the mean functions, unless that's wrong...

[csuter]

AFAIK, the only way to do posterior sampling of other model parameters would be MCMC. The GP regression model docstring has an example of doing sampling of kernel params (and IIRC, maybe the observation noise?). You could try something similar for the mean params.

[csuter]

(Search for "Marginalization of model hyperparameters" on https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/GaussianProcessRegressionModel)

[brianwa84]

Reopen if you're still having issues here