Bayesian model calibration with GP in PyMC3

I am attempting to implement bayesian model calibration under the classical Kennedy-O'Hagan framework using PyMC3. Briefly, the setup is as follows:

I have $n$ observations $z_i$ of the form $$z_i = f(x_i, \theta) + g(x_i) + e_i$$, where $x_i$ are known imputs and $\theta$ are unknown calibration parameters and $e_i$ are iid error terms. We also have $m$ model evaluations $y_j$ of the form $$y_j = f(x__j, \theta_j)$$, where both $x_j$ and $\theta_j$ are known. Threfore, the data consists of all $z_i$ and $y_j$. In the paper, Kennedy-O'Hagan model f, g using gaussian processes:

$$f \sim GP\{m(.,.), \Sigma\[(.,.),(.,.)\] \}$$
$$g \sim GP\{m(.), \Sigma\[(.),(.)\] \}$$.

Among other things, the goal is to get posterior samples for the unknow calibration imputs $\theta$. As I said, I am trying to implement the solution using Python and PyMC3 in particular.

I have the following challenges:

  • How to define the model in PyMC3 in the first place? I know how to do it ,when I have one GP for the whole data set or sum of GP for the whole data set. In this case, part of the data follows one model, another part follows another model.

  • $\theta$ is an input into the covariance function for f. There does not seem any streight forward way to do this in PyMC3 as well.

I will be very grateful for any tips or ideas how to solve the above mentioned problems!