# Gibbs Sampling
This tutorial extends the [linear regression
tutorial](01a-lin-reg.md#linear-regression). Here, we show how to sample
model parameters using a Gibbs kernel.
As this tutorial is a continuation of the previous tutorials, we will
use the same model and data assumed there.
## Data and imports
``` python
import jax
import jax.numpy as jnp
import numpy as np
# We use distributions and bijectors from tensorflow probability
import tensorflow_probability.substrates.jax.distributions as tfd
import tensorflow_probability.substrates.jax.bijectors as tfb
import liesel.goose as gs
import liesel.model as lsl
import matplotlib.pyplot as plt
```
``` python
# Generate data
rng = np.random.default_rng(42)
# sample size and true parameters
n = 500
true_beta = np.array([1.0, 2.0])
true_sigma = 1.0
# data-generating process
x0 = rng.uniform(size=n)
X_mat = np.column_stack([np.ones(n), x0])
eps = rng.normal(scale=true_sigma, size=n)
y_vec = X_mat @ true_beta + eps
# define beta
beta_prior = lsl.Dist(tfd.Normal, loc=0.0, scale=100.0)
beta = lsl.Var.new_param(value=jnp.array([0.0, 0.0]), dist=beta_prior, name="beta")
# define the variance and the scale
a = lsl.Var.new_param(0.01, name="a")
b = lsl.Var.new_param(0.01, name="b")
sigma_sq_prior = lsl.Dist(tfd.InverseGamma, concentration=a, scale=b)
sigma_sq = lsl.Var.new_param(value=1.0, dist=sigma_sq_prior, name="sigma_sq")
# Define sigma as a transformation of sigma_sq for the likelihood
sigma = lsl.Var.new_calc(jnp.sqrt, sigma_sq, name="sigma")
# calculator-setup
X = lsl.Var.new_obs(X_mat, name="X")
mu = lsl.Var.new_calc(jnp.dot, X, beta, name="mu")
# Build response
y_dist = lsl.Dist(tfd.Normal, loc=mu, scale=sigma)
y = lsl.Var.new_obs(y_vec, dist=y_dist, name="y")
# Plot model
model = lsl.Model([y])
model.plot()
```
## MCMC inference
### Using a Gibbs kernel
This time we want to sample the previously fixed `sigma_sq` with a Gibbs
sampler. Using a Gibbs kernel is a bit more complicated, because Goose
doesn’t automatically derive the full conditional from the model graph.
Hence, the user needs to provide a function to sample from the full
conditional. The function needs to accept a PRNG key and a model state
as arguments, and it needs to return a dictionary with the variable name
as the key and the new variable value as the value. We could also update
multiple parameters with one Gibbs kernel by returning a dictionary with
several entries.
For this normal-inverse-gamma model, the full conditional of $\sigma^2$
is again an inverse-gamma distribution. To retrieve the relevant values
from the `model_state`, we use {meth}`.Model.extract_position`.
``` python
def draw_sigma_sq(prng_key, model_state):
# extract relevant values from model state
pos = model.extract_position(
position_keys=["y", "mu", "sigma_sq", "a", "b"], model_state=model_state
)
# calculate relevant intermediate quantities
n = len(pos["y"])
resid = pos["y"] - pos["mu"]
a_gibbs = pos["a"] + n / 2
b_gibbs = pos["b"] + jnp.sum(resid**2) / 2
# draw new value from full conditional
draw = b_gibbs / jax.random.gamma(prng_key, a_gibbs)
# return key-value pair of variable name and new value
return {"sigma_sq": draw}
```
The regression coefficients `beta` are still sampled with NUTS. For
`sigma_sq`, we attach an {class}`~.goose.MCMCSpec` with
{meth}`~.goose.GibbsKernel.with_transition_fn`, which turns our custom
transition function into a kernel factory that {class}`.LieselMCMC` can
use. The Gibbs kernel itself does not need adaptation, but the NUTS
kernel for `beta` does, so we still run an adaptation phase before
drawing posterior samples.
``` python
beta.inference = gs.MCMCSpec(gs.NUTSKernel)
sigma_sq.inference = gs.MCMCSpec(gs.GibbsKernel.with_transition_fn(draw_sigma_sq))
results = gs.LieselMCMC(model).run_for_epochs(
seed=1, num_chains=4, adaptation=1000, posterior=1000
)
```
liesel.goose.mcmc_spec - WARNING - No inference specification defined for Var(name="b"). If you do not add a kernel for this parameter manually to an EngineBuilder, it will not be sampled.
liesel.goose.mcmc_spec - WARNING - No inference specification defined for Var(name="a"). If you do not add a kernel for this parameter manually to an EngineBuilder, it will not be sampled.
liesel.goose.builder - WARNING - No jitter functions provided for position keys 'sigma_sq', 'beta'. The initial values for these keys won't be jittered
liesel.goose.engine - INFO - Initializing kernels...
liesel.goose.engine - INFO - Done
liesel.goose.engine - INFO - Starting epoch: FAST_ADAPTATION, 100 transitions, 25 jitted together
0%| | 0/4 [00:00
Parameter summary:
|
|
|
kernel
|
mean
|
sd
|
q_0.05
|
q_0.5
|
q_0.95
|
sample_size
|
ess_bulk
|
ess_tail
|
rhat
|
|
parameter
|
index
|
|
|
|
|
|
|
|
|
|
|
|
beta
|
(0,)
|
kernel_01
|
0.983
|
0.093
|
0.832
|
0.983
|
1.133
|
4000
|
1063.525
|
1171.567
|
1.005
|
|
(1,)
|
kernel_01
|
1.912
|
0.159
|
1.652
|
1.912
|
2.172
|
4000
|
1009.363
|
1118.012
|
1.007
|
|
sigma_sq
|
()
|
kernel_00
|
1.043
|
0.066
|
0.939
|
1.041
|
1.154
|
4000
|
4090.345
|
4014.613
|
1.000
|
Acceptance probabilities:
|
|
|
|
acceptance_probability
|
position_moved
|
|
kernel
|
positions
|
phase
|
|
|
|
kernel_00
|
sigma_sq
|
posterior
|
1.000
|
1.000
|
|
warmup
|
1.000
|
1.000
|
|
kernel_01
|
beta
|
posterior
|
0.887
|
NaN
|
|
warmup
|
0.791
|
NaN
|
Error summary:
|
|
|
|
|
|
count
|
sample_size
|
sample_size_total
|
relative
|
|
kernel
|
positions
|
error_code
|
error_msg
|
phase
|
|
|
|
|
|
kernel_01
|
beta
|
1
|
divergent transition
|
warmup
|
51
|
4000
|
4000
|
0.013
|
|
posterior
|
0
|
4000
|
4000
|
0.000
|
And plot them.
``` python
gs.plot_trace(results)
```
