Migrating to PyMC4 from PyMC3 Multidimensional GP (original) (raw)

Dear all,
My apologies if the question is rudimentary. I am following a notebook from @fonnesbeck presented at PyData NYC 2019, shared here on multidimensional GP, I am using the newly released PyMC4, I am not able to do any MCMC sampling or even doing any pm.find_MAP(). I keep getting an error which I think is related to the specification of sf2 variable. The error reads " ** ValueError : Random variables detected in the logp graph: [l]. This can happen when DensityDist logp or Interval transform functions reference nonlocal variables.**"

nd = 15
xu1, xu2 = np.meshgrid(np.linspace(0, 300, nd), np.linspace(0, 300, nd))
Xu = np.concatenate([xu1.reshape(nd*nd, 1), xu2.reshape(nd*nd, 1)], 1)

with pm.Model() as spatial_model:
    
    l = pm.HalfCauchy("l", beta=3, shape=(2,))
    sf2 = pm.HalfCauchy("sf2", beta=3)
    sn2 = pm.HalfCauchy("sn2", beta=3)

    K = pm.gp.cov.ExpQuad(2, l) * sf2**2
    
    gp_spatial = pm.gp.MarginalSparse(cov_func=K, approx="FITC")
    obs = gp_spatial.marginal_likelihood("obs", X=X_obs, Xu=Xu, y=y_obs, noise=sn2)

    mp = pm.find_MAP()


nd = 30
z1, z2 = np.meshgrid(np.linspace(0, 300, nd), np.linspace(0, 300, nd))
Z = np.concatenate([z1.reshape(nd*nd, 1), z2.reshape(nd*nd, 1)], 1)


with spatial_model:

    f_pred = gp_spatial.conditional('f_pred', Z)
    
    samples = pm.sample_posterior_predictive([mp], vars=[f_pred], samples=100)

with sns.axes_style("white"):

    plt.figure(figsize=(10,8))
    ax = sns.heatmap(samples['f_pred'].mean(0).reshape(nd, nd), cmap='viridis')
    ax.invert_yaxis()
    ax.set_yticklabels([])
    ax.set_xticklabels([])

My questions are:

  1. Is there any thing that I could change in the model so that it could carry on with the sampling given that I am using PyMC4?
  2. Is there anything wrong with obs = gp_spatial.marginal_likelihood("obs", X=X_obs, Xu=Xu, y=y_obs, noise=sn2) that has to be updated to PyMC4 syntax?
  3. How about changes in PyMC4 calling the sample_posterior_predictive, specifically do I need to change anything in the lines
   f_pred = gp_spatial.conditional('f_pred', Z)
   
   samples = pm.sample_posterior_predictive([mp], vars=[f_pred], samples=100)

Many thanks,
Ali