# Comparing samplers In this tutorial, we compare two different sampling schemes on the `mcycle` dataset with a Gaussian location-scale regression model and two splines for the mean and the standard deviation. The `mcycle` dataset is a “data frame giving a series of measurements of head acceleration in a simulated motorcycle accident, used to test crash helmets” (from the help page). It contains the following two variables: - `times`: in milliseconds after impact - `accel`: in g We set up the model in Python with [Liesel-GAM](https://github.com/liesel-devs/liesel_gam), using {class}`liesel_gam.TermBuilder` for the P-spline terms. See the [Liesel-GAM documentation and examples](https://github.com/liesel-devs/liesel_gam#readme) for more information about additive terms and predictors. We load the data set from R with [ryp](https://github.com/Wainberg/ryp) and then continue with a pure Python model specification and sampling workflow. ``` python from pathlib import Path import jax.numpy as jnp import matplotlib.pyplot as plt import numpy as np import seaborn as sns import tensorflow_probability.substrates.jax.bijectors as tfb import tensorflow_probability.substrates.jax.distributions as tfd import liesel.goose as gs import liesel.model as lsl import liesel_gam as gam from ryp import r, to_py ``` We start by loading the data set from the R package `MASS` and converting it to a pandas data frame. ``` python r("library(MASS)") r("data(mcycle); mcycle <- as.data.frame(mcycle)") mcycle = to_py("mcycle", format="pandas") ``` ``` python fig, ax = plt.subplots(figsize=(8, 4)) sns.scatterplot(data=mcycle, x="times", y="accel", color="0.25", s=35, ax=ax) ax.set(xlabel="time after impact", ylabel="acceleration", title="mcycle data") plt.show() ``` Next, we build the Gaussian location-scale model. Both distributional parameters use an additive predictor with an intercept and a P-spline in `times`. The scale predictor uses an exponential inverse link to keep the standard deviation positive. The `TermBuilder` is initialized with an IWLS MCMC specification, so the P-spline regression coefficients are sampled with IWLS kernels by default. The additive predictor intercepts also use their default IWLS inference specification. The smoothing variances of the P-splines receive Gibbs kernels automatically. ``` python tb = gam.TermBuilder.from_df(mcycle, default_inference=gs.MCMCSpec(gs.IWLSKernel)) loc = gam.AdditivePredictor("loc") scale = gam.AdditivePredictor("scale", inv_link=jnp.exp) loc_smooth = tb.ps("times", k=20, prefix="loc.") scale_smooth = tb.ps("times", k=20, prefix="scale.") loc += loc_smooth scale += scale_smooth response_dist = lsl.Dist(tfd.Normal, loc=loc, scale=scale) y = lsl.Var.new_obs(mcycle["accel"], response_dist, name="y") model = lsl.Model(y) ``` ## Metropolis-in-Gibbs First, we run the model with the inference specifications attached during model construction. This gives a Metropolis-in-Gibbs sampling scheme with IWLS kernels for the regression coefficients ($\boldsymbol{\beta}$) and Gibbs kernels for the smoothing variances ($\tau^2$) of the splines. ``` python iwls_results = gs.LieselMCMC(model).run_for_epochs( seed=1, num_chains=4, adaptation=1000, posterior=10000, posterior_thinning=10, show_progress=False ) ``` liesel.goose.builder - WARNING - No jitter functions provided for position keys '$\\beta_{loc.ps(times)}$', '$\\beta_{scale.ps(times)}$', '$\\tau_{loc.ps(times)}^2$', '$\\beta_{0,loc}$', '$\\tau_{scale.ps(times)}^2$', '$\\beta_{0,scale}$'. 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 - Finished warmup Clearly, the performance of the sampler could be better, especially for the intercept of the mean. The corresponding chain exhibits a very strong autocorrelation. ``` python gs.Summary(iwls_results) ```

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,loc}$ () kernel_03 -25.231 1.982 -28.440 -25.266 -21.932 4000 91.735 99.749 1.043
$\beta_{0,scale}$ () kernel_05 2.724 0.075 2.601 2.723 2.853 4000 1240.162 2309.053 1.002
$\beta_{loc.ps(times)}$ (0,) kernel_00 2.389 10.804 -15.710 2.433 19.814 4000 2745.445 3217.044 1.000
(1,) kernel_00 11.079 10.199 -4.826 10.582 28.691 4000 2502.842 3212.205 1.000
(2,) kernel_00 1.566 9.327 -13.700 1.507 17.078 4000 2734.069 3514.793 1.000
(3,) kernel_00 -3.259 8.896 -18.092 -3.024 10.990 4000 3171.874 3456.306 1.001
(4,) kernel_00 -9.310 8.856 -24.442 -9.072 4.528 4000 3452.395 3637.538 1.001
(5,) kernel_00 -10.120 8.434 -24.343 -9.852 3.212 4000 2684.143 3569.892 1.001
(6,) kernel_00 0.110 7.682 -12.099 -0.028 12.989 4000 2973.361 2888.397 1.003
(7,) kernel_00 -0.425 7.055 -12.029 -0.529 11.258 4000 3317.425 3522.206 1.000
(8,) kernel_00 10.682 6.637 -0.434 10.657 21.660 4000 3410.736 3716.644 1.000
(9,) kernel_00 -15.615 5.729 -25.140 -15.519 -6.263 4000 2888.242 2849.860 1.001
(10,) kernel_00 7.597 4.886 -0.181 7.516 15.832 4000 2699.041 3273.765 1.001
(11,) kernel_00 -23.715 4.428 -31.150 -23.687 -16.500 4000 3119.502 3535.451 1.000
(12,) kernel_00 9.336 3.280 4.172 9.334 14.651 4000 3325.456 3589.285 1.001
(13,) kernel_00 -10.207 2.583 -14.431 -10.168 -6.068 4000 3216.028 3244.630 1.001
(14,) kernel_00 12.242 1.877 9.184 12.259 15.173 4000 3356.569 3364.785 1.001
(15,) kernel_00 2.242 1.235 0.189 2.249 4.188 4000 1827.116 2513.063 1.003
(16,) kernel_00 -3.139 0.646 -4.195 -3.124 -2.139 4000 3279.625 3283.467 1.001
(17,) kernel_00 0.912 0.242 0.524 0.919 1.291 4000 943.454 2046.986 1.004
(18,) kernel_00 2.983 0.925 1.523 3.007 4.375 4000 3234.198 2806.174 1.002
$\beta_{scale.ps(times)}$ (0,) kernel_01 0.011 0.145 -0.213 0.005 0.256 4000 1157.099 1448.394 1.005
(1,) kernel_01 -0.024 0.149 -0.280 -0.016 0.198 4000 1200.021 1000.396 1.005
(2,) kernel_01 0.006 0.142 -0.215 0.002 0.238 4000 1351.423 1157.364 1.005
(3,) kernel_01 0.026 0.142 -0.194 0.022 0.264 4000 1183.395 1030.918 1.006
(4,) kernel_01 0.031 0.146 -0.191 0.022 0.284 4000 1099.421 722.438 1.007
(5,) kernel_01 -0.033 0.144 -0.272 -0.025 0.191 4000 1001.900 985.622 1.005
(6,) kernel_01 -0.056 0.145 -0.309 -0.043 0.159 4000 1086.999 1004.761 1.003
(7,) kernel_01 -0.017 0.133 -0.238 -0.015 0.196 4000 1058.990 1067.830 1.007
(8,) kernel_01 0.103 0.147 -0.100 0.080 0.374 4000 664.293 1027.972 1.011
(9,) kernel_01 -0.094 0.133 -0.324 -0.081 0.095 4000 870.111 1235.133 1.005
(10,) kernel_01 -0.116 0.120 -0.323 -0.108 0.059 4000 1030.726 1390.782 1.004
(11,) kernel_01 0.008 0.112 -0.169 0.011 0.185 4000 1044.794 1269.849 1.004
(12,) kernel_01 0.202 0.126 0.027 0.185 0.432 4000 435.379 606.760 1.011
(13,) kernel_01 0.138 0.099 -0.009 0.131 0.306 4000 625.898 1139.024 1.003
(14,) kernel_01 -0.079 0.085 -0.228 -0.072 0.046 4000 569.673 851.310 1.006
(15,) kernel_01 0.044 0.056 -0.039 0.039 0.138 4000 676.537 1207.682 1.005
(16,) kernel_01 0.026 0.032 -0.027 0.028 0.074 4000 582.852 778.132 1.014
(17,) kernel_01 -0.063 0.012 -0.081 -0.064 -0.042 4000 657.701 1249.209 1.007
(18,) kernel_01 0.111 0.048 0.033 0.111 0.191 4000 735.290 1102.061 1.013
$\tau_{loc.ps(times)}^2$ () kernel_02 137.231 66.379 61.990 122.915 262.839 4000 2799.840 3249.357 1.000
$\tau_{scale.ps(times)}^2$ () kernel_04 0.022 0.021 0.003 0.016 0.060 4000 303.019 492.178 1.021

Acceptance probabilities:

acceptance_probability position_moved
kernel positions phase
kernel_00 $\beta_{loc.ps(times)}$ posterior 0.863 0.863
warmup 0.794 0.793
kernel_01 $\beta_{scale.ps(times)}$ posterior 0.847 0.848
warmup 0.793 0.790
kernel_02 $\tau_{loc.ps(times)}^2$ posterior 1.000 1.000
warmup 1.000 1.000
kernel_03 $\beta_{0,loc}$ posterior 0.922 0.923
warmup 0.923 0.927
kernel_04 $\tau_{scale.ps(times)}^2$ posterior 1.000 1.000
warmup 1.000 1.000
kernel_05 $\beta_{0,scale}$ posterior 0.906 0.905
warmup 0.912 0.909
``` python gs.plot_trace(iwls_results) ``` To confirm that the chains have converged to reasonable values, we plot the posterior mean of the location predictor together with a 90% credible interval: ``` python def plot_loc_estimate(results, model, title): samples = results.get_posterior_samples() loc_samples = model.vars["loc"].predict(samples) loc_summary = gs.SamplesSummary.from_array( loc_samples, name="loc", which=["mean", "quantiles"], ) loc_summary_df = loc_summary.to_dataframe().reset_index() loc_summary_df["times"] = mcycle["times"].to_numpy() plot_data = ( loc_summary_df[["times", "mean", "q_0.05", "q_0.95"]] .groupby("times", as_index=False) .mean() .sort_values("times") ) fig, ax = plt.subplots(figsize=(8, 5)) ax.fill_between( plot_data["times"], plot_data["q_0.05"], plot_data["q_0.95"], color=sns.color_palette()[1], alpha=0.25, label="90% credible interval", ) sns.lineplot( data=plot_data, x="times", y="mean", color=sns.color_palette()[1], linewidth=2, label="posterior mean", ax=ax, ) sns.scatterplot( data=mcycle, x="times", y="accel", color="0.25", s=25, ax=ax, label="observed data", ) ax.set(xlabel="time after impact", ylabel="acceleration", title=title) plt.show() plot_loc_estimate(iwls_results, model, "Estimated mean function (IWLS/Gibbs)") ``` ## NUTS sampler As an alternative, we use NUTS kernels for the spline-specific parameter blocks. The helper below copies the model graph, log-transforms the smoothing variances by bijecting them with an exponential bijector, and assigns one NUTS kernel group per additive term. ``` python def strategy_term_blocked( model: lsl.Model, predictors: list[str], kernel_constructor, **kwargs ): model = model.copy() for k, v in model.parameters.items(): if "tau" in k: v.biject(tfb.Exp(), inference="drop") for predictor_name in predictors: predictor = model.vars[predictor_name] if predictor.intercept: predictor.intercept.inference = gs.MCMCSpec( kernel_constructor, kernel_kwargs=kwargs ) for term in predictor.terms.values(): for param in model.parental_submodel(term).parameters.values(): model.parameters[param.name].inference = gs.MCMCSpec( kernel_constructor, kernel_group=term.name, kernel_kwargs=kwargs ) return model ``` ``` python nuts_model = strategy_term_blocked(model, ["loc", "scale"], gs.NUTSKernel) ``` The resulting model contains transformed smoothing variances on the unconstrained log scale. Here is the transformed model graph: ``` python nuts_model.plot() ``` Now we can run the sampler from the `MCMCSpec` objects stored in the model. In complex models like this one, it can be beneficial to sample the parameters of each additive term in a separate NUTS block. ``` python nuts_results = gs.LieselMCMC(nuts_model).run_for_epochs( seed=1, num_chains=4, adaptation=1000, posterior=1000, show_progress=False ) ``` liesel.goose.builder - WARNING - No jitter functions provided for position keys '$\\beta_{loc.ps(times)}$', 'h($\\tau_{loc.ps(times)}^2$)', '$\\beta_{scale.ps(times)}$', 'h($\\tau_{scale.ps(times)}^2$)', '$\\beta_{0,loc}$', '$\\beta_{0,scale}$'. 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 - Finished warmup The blocked NUTS strategy overall seems to do a good job and can yield higher effective sample sizes than the IWLS sampler, especially for the spline coefficients of the scale model. ``` python gs.Summary(nuts_results) ```

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,loc}$ () kernel_02 -25.157 2.255 -29.277 -25.095 -21.618 4000 13.116 15.298 1.240
$\beta_{0,scale}$ () kernel_03 2.723 0.076 2.603 2.720 2.851 4000 675.615 1260.077 1.004
$\beta_{loc.ps(times)}$ (0,) kernel_00 2.682 10.627 -14.695 2.687 19.847 4000 4064.084 2753.618 1.001
(1,) kernel_00 11.372 10.339 -4.712 10.733 29.440 4000 2718.144 2039.869 1.003
(2,) kernel_00 1.461 9.325 -14.031 1.412 16.753 4000 3319.317 2426.635 1.002
(3,) kernel_00 -3.353 8.916 -17.948 -3.249 11.335 4000 3185.618 2749.068 1.002
(4,) kernel_00 -9.362 8.880 -24.667 -9.085 4.699 4000 3058.779 2325.766 1.001
(5,) kernel_00 -10.313 8.543 -24.531 -10.120 3.330 4000 2856.079 2586.780 1.001
(6,) kernel_00 0.566 7.934 -12.275 0.369 13.514 4000 2530.758 2844.444 1.002
(7,) kernel_00 -0.235 7.445 -12.426 -0.319 12.225 4000 2481.104 2443.635 1.001
(8,) kernel_00 10.731 6.754 -0.296 10.691 22.074 4000 2269.861 2242.519 1.002
(9,) kernel_00 -15.589 5.883 -25.395 -15.445 -6.031 4000 1590.324 2105.093 1.003
(10,) kernel_00 7.504 4.913 -0.469 7.489 15.721 4000 1014.294 1676.991 1.004
(11,) kernel_00 -23.858 4.530 -31.420 -23.791 -16.633 4000 1352.847 1894.121 1.002
(12,) kernel_00 9.164 3.213 3.938 9.161 14.466 4000 1170.496 1725.139 1.006
(13,) kernel_00 -10.156 2.535 -14.347 -10.178 -6.043 4000 1191.233 1847.846 1.005
(14,) kernel_00 12.291 1.875 9.226 12.301 15.349 4000 1192.178 1541.098 1.004
(15,) kernel_00 2.294 1.234 0.292 2.276 4.301 4000 522.760 1089.537 1.012
(16,) kernel_00 -3.120 0.627 -4.142 -3.120 -2.098 4000 989.886 1481.136 1.005
(17,) kernel_00 0.914 0.246 0.510 0.917 1.307 4000 96.340 863.558 1.037
(18,) kernel_00 3.011 0.907 1.557 3.008 4.498 4000 928.757 1363.198 1.004
$\beta_{scale.ps(times)}$ (0,) kernel_01 0.001 0.141 -0.219 0.001 0.217 4000 4692.991 1962.539 1.002
(1,) kernel_01 -0.024 0.143 -0.273 -0.016 0.192 4000 4934.493 1801.269 1.003
(2,) kernel_01 0.011 0.146 -0.211 0.007 0.251 4000 5073.525 1989.663 1.005
(3,) kernel_01 0.029 0.142 -0.186 0.022 0.264 4000 4361.999 1997.520 1.001
(4,) kernel_01 0.026 0.144 -0.192 0.020 0.257 4000 4809.598 1909.172 1.003
(5,) kernel_01 -0.032 0.144 -0.281 -0.026 0.188 4000 5427.737 1621.686 1.003
(6,) kernel_01 -0.056 0.143 -0.302 -0.046 0.156 4000 3333.762 1551.118 1.000
(7,) kernel_01 -0.012 0.132 -0.225 -0.009 0.196 4000 4524.246 2200.880 1.003
(8,) kernel_01 0.093 0.145 -0.109 0.075 0.361 4000 1907.307 1375.404 1.002
(9,) kernel_01 -0.090 0.125 -0.308 -0.076 0.090 4000 2468.683 2047.995 1.003
(10,) kernel_01 -0.120 0.122 -0.337 -0.107 0.057 4000 2564.739 2204.354 1.002
(11,) kernel_01 0.009 0.114 -0.181 0.009 0.195 4000 3307.058 2243.883 1.002
(12,) kernel_01 0.204 0.122 0.028 0.193 0.421 4000 814.724 1340.109 1.003
(13,) kernel_01 0.138 0.100 -0.009 0.133 0.309 4000 1246.340 1565.672 1.002
(14,) kernel_01 -0.083 0.085 -0.231 -0.078 0.043 4000 971.161 1790.243 1.002
(15,) kernel_01 0.044 0.058 -0.044 0.040 0.145 4000 1184.525 1425.558 1.004
(16,) kernel_01 0.024 0.033 -0.033 0.026 0.074 4000 1052.191 1817.478 1.004
(17,) kernel_01 -0.063 0.013 -0.083 -0.063 -0.040 4000 1309.675 1534.598 1.001
(18,) kernel_01 0.108 0.049 0.026 0.109 0.187 4000 1391.988 1917.419 1.004
h($\tau_{loc.ps(times)}^2$) () kernel_00 4.835 0.440 4.135 4.827 5.593 4000 1654.489 1923.113 1.001
h($\tau_{scale.ps(times)}^2$) () kernel_01 -4.184 0.841 -5.620 -4.138 -2.859 4000 363.985 693.435 1.006

Acceptance probabilities:

acceptance_probability position_moved
kernel positions phase
kernel_00 $\beta_{loc.ps(times)}$, h($\tau_{loc.ps(times)}^2$) posterior 0.892 NaN
warmup 0.793 NaN
kernel_01 $\beta_{scale.ps(times)}$, h($\tau_{scale.ps(times)}^2$) posterior 0.876 NaN
warmup 0.794 NaN
kernel_02 $\beta_{0,loc}$ posterior 0.862 NaN
warmup 0.791 NaN
kernel_03 $\beta_{0,scale}$ posterior 0.875 NaN
warmup 0.793 NaN

Error summary:

count sample_size sample_size_total relative
kernel positions error_code error_msg phase
kernel_00 $\beta_{loc.ps(times)}$, h($\tau_{loc.ps(times)}^2$) 1 divergent transition warmup 364 4000 4000 0.091
posterior 18 4000 4000 0.004
2 maximum tree depth warmup 393 4000 4000 0.098
posterior 0 4000 4000 0.000
kernel_01 $\beta_{scale.ps(times)}$, h($\tau_{scale.ps(times)}^2$) 1 divergent transition warmup 272 4000 4000 0.068
posterior 0 4000 4000 0.000
2 maximum tree depth warmup 8 4000 4000 0.002
posterior 0 4000 4000 0.000
kernel_02 $\beta_{0,loc}$ 1 divergent transition warmup 75 4000 4000 0.019
posterior 0 4000 4000 0.000
kernel_03 $\beta_{0,scale}$ 1 divergent transition warmup 36 4000 4000 0.009
posterior 0 4000 4000 0.000
``` python gs.plot_trace(nuts_results) ``` Again, here is the posterior mean function with a 90% credible interval: ``` python plot_loc_estimate(nuts_results, nuts_model, "Estimated mean function (NUTS)") ```