Interleave sampling of two groups of variables in pymc3

I am trying to experiment with fully Bayesian inference in GP classification. The use case is to draw samples of the hypers (theta) and latent f jointly. I want to implement a scheme where the sampling is interleaved so the pseudo code looks something like this:

Step 1-> Sample some hypers, freeze those
Step 2-> Sample from the conditional p( f | theta, y).

The idea is to navigate in the constrained f space (high dimensional) using different rules than the low dimensional theta space.

I wrote some trial code using the compound step but not sure if it is giving me f samples from the conditional.

HMC Model

    with pm.Model() as model:

        # covariance function
        log_l = pm.Normal('log_l', mu=0, sigma=3)
        log_s = pm.Normal('log_s', mu=0, sigma=3)
     
        l = pm.Deterministic('l', pm.math.exp(log_l))
        s = pm.Deterministic('s', pm.math.exp(log_s))
       
        cov = s**2 * pm.gp.cov.ExpQuad(1, l)

        gp = pm.gp.Latent(cov_func=cov)

        # make gp prior
        f = gp.prior("f", X=x[:,None], reparameterize=True)
        
        # logit link and Bernoulli likelihood
        p = pm.Deterministic("p", pm.math.invlogit(f))
        y_ = pm.Bernoulli("y", p=p, observed=y)

        step_theta = pm.Metropolis([log_l, log_s])
        step_latent = pm.NUTS([f])
        trace_hmc = pm.sample(draws=100, tune=200, chains=1, step=[step_theta, step_latent])
1 Like

Yes, what you are doing is correct. The compound step will first sample log_l and log_s using Metropolis, conditioned on the value of f from the previous sample, and then sample f using NUTS conditioned on the log_l and log_s.

1 Like

Thanks for writing back, so my intention is to intervleave the sampling between variables using the same sampler but different parameters: for instance, HMC with different path lengths for theta and f like below:

  # HMC Model 
    with pm.Model() as model:
    
            # covariance function
            log_l = pm.Normal('log_l', mu=0, sigma=3)
            log_s = pm.Normal('log_s', mu=0, sigma=3)
         
            l = pm.Deterministic('l', pm.math.exp(log_l))
            s = pm.Deterministic('s', pm.math.exp(log_s))
           
            cov = s**2 * pm.gp.cov.ExpQuad(1, l)
    
            gp = pm.gp.Latent(cov_func=cov)
    
            # make gp prior
            f = gp.prior("f", X=X_100[:,None], reparameterize=True)
            
            # logit link and Bernoulli likelihood
            p = pm.Deterministic("p", pm.math.invlogit(f))
            y_ = pm.Bernoulli("y", p=p, observed=y_100)
    
            step_theta = pm.step_methods.HamiltonianMC(vars=[log_l, log_s], is_cov=True)
            step_latent = pm.step_methods.HamiltonianMC(vars=[f], path_length=10)
            trace_hmc_100 = pm.sample(draws=200, tune=200, chains=1, step=[step_theta, step_latent])

But I am not sure why the logs shown during sampling read all the three variables for the second step ?

CompoundStep
HamiltonianMC: [log_s, log_l]
HamiltonianMC: [log_l, log_s, f_rotated_]
Sampling chain 0, 0 divergences: 12%|█▏ | 47/400 [00:08<01:03, 5.53it/s]

Ah I see - the vars that step methods recognized are the variables in the latent space - try

step_latent = pm.step_methods.HamiltonianMC(vars=[model.named_vars['f_rotated_']], path_length=10)
1 Like

that does it!!

Sequential sampling (1 chains in 1 job)
CompoundStep
>NUTS: [log_s, log_l]
>NUTS: [f_rotated_]
Sampling chain 0, 0 divergences: 100%|██████████| 600/600 [00:04<00:00, 132.69it/s]