# MCMC sampling

Whittaker's default inference pipeline is the **Laplace approximation**: after P-IRLS finds the posterior mode \hat\beta, the posterior is approximated as q(\beta) = \mathcal{N}(\hat\beta,\\ V\_\beta), where V\_\beta = \phi\\(X^\top W X + \sum_j \lambda_j S_j)^{-1}. For Gaussian response this is exact; for non-Gaussian families it can underestimate spread when the likelihood surface is asymmetric.

**MCMC** (via the No-U-Turn Sampler, NUTS) draws directly from the exact posterior, so no Gaussian approximation is made. NUTS uses the gradient of the log-posterior to grow a binary trajectory tree in both directions, stopping automatically when the path would double back on itself. This gives it the efficiency of Hamiltonian Monte Carlo without requiring you to tune a trajectory length. The result is a collection of coefficient vectors that represent the full posterior distribution of the smooth functions.


# When to use MCMC

Use `method="MCMC"` when:

- you need the most accurate posterior uncertainty for non-Gaussian families (Poisson, Binomial, Gamma) and neither the Laplace approximation nor VI is sufficient
- the posterior may be multimodal or strongly skewed
- you want posterior predictive distributions for individual observations, not just marginal standard errors
- you are doing formal Bayesian inference and need to verify convergence via R-hat and ESS

For large n or when a fast answer is needed, the Laplace approximation (`method="REML"`, the default) is usually adequate. For non-Gaussian families with moderate n, variational inference (`method="VI"`) offers a nice middle ground: better calibrated than Laplace, cheaper than MCMC.


# Basic usage

Pass `method="MCMC"` to `fit()`. Everything else stays the same.


``` python
import numpy as np
import whittaker as wk
from whittaker.families.poisson import Poisson

rng = np.random.default_rng(23)
n = 300
x = np.linspace(0, 2 * np.pi, n)
lam = np.exp(1.5 * np.sin(x))
y = rng.poisson(lam).astype(float)

data = {"x": x, "y": y}

model = wk.GAM("y ~ s(x)", family=Poisson())
model.fit(data, method="MCMC")
model.summary()
```


    GAM fit summary
    ============================================================
    Formula:    y ~ s(x)
    Family:     Poisson(link='log')
    Inference:  MCMC (NUTS, 4 chains × 1000 draws)
    Observations: 300
    Coefficients: 10

    Total EDF:  6.79
    Scale est:  1.000000
    Draws:      4 chains × 1000 = 4000
    Warmup:     500 per chain
    Acceptance: 0.813
    Step size:  0.01643
    Max R-hat:  1.0023
    Min ESS (bulk): 1322
    Min ESS (tail): 1587
    Tree depth: 7.84 (mean)
    Deviance:   312.0481
    Null dev:   742.3931
    Dev. expl:  58.0%
    AIC:        808.60
    BIC:        833.75


The summary shows `Inference: MCMC (NUTS, …)` along with acceptance rate, R-hat, ESS, mean tree depth, and -- if any occur -- a divergent-transition warning. GCV, AIC, and BIC are not defined for MCMC fits.


# Inspecting the MCMC result

The [mcmc_result](../reference/GAM.md#whittaker.GAM.mcmc_result) property returns an `MCMCResult` object with the full posterior sample and convergence diagnostics:


``` python
mr = model.mcmc_result

print(f"Samples shape:    {mr.samples.shape}")  # (p, n_chains * n_samples)
print(f"R-hat (max):      {mr.r_hat.max():.4f}")
print(f"ESS bulk (min):   {mr.ess.min():.1f}")
print(f"ESS tail (min):   {mr.ess_tail.min():.1f}")
print(f"Acceptance rate:  {mr.acceptance_rate:.3f}")  # mean per-leaf α for NUTS
print(f"Mean tree depth:  {mr.mean_tree_depth:.2f}")  # NUTS only
print(f"Divergences:      {mr.n_divergent}")
```


    Samples shape:    (10, 4000)
    R-hat (max):      1.0023
    ESS bulk (min):   1322.3
    ESS tail (min):   1587.0
    Acceptance rate:  0.813
    Mean tree depth:  7.84
    Divergences:      0


[mcmc_result](../reference/GAM.md#whittaker.GAM.mcmc_result) is `None` if the model was fitted with any other method.


## R-hat and ESS

**R-hat** is the rank-normalized split R-hat (Vehtari et al. 2021). Each chain is first split in half, giving twice as many half-chains. Classic Gelman-Rubin R-hat is then applied to the rank-normalized draws. Splitting detects non-stationarity within a single chain. Rank normalization makes the statistic robust to heavy-tailed posteriors. Values below 1.01 are ideal. Values below 1.1 are generally acceptable; values above 1.1 suggest insufficient warmup, too few chains, or poor geometry.

**ESS bulk** (`ess`) measures mixing in the bulk of the posterior. It applies the standard autocorrelation ESS estimator to the rank-normalized draws. An ESS ratio (ESS / total draws) above 0.05 is generally adequate. The value below 0.05 suggests the chains are heavily autocorrelated.

**ESS tail** (`ess_tail`) measures how reliably the sampler visits the tails. It is the minimum of the ESS of the binary indicator I(x \le Q\_{0.05}) and I(x \ge Q\_{0.95}) across all draws. Low tail ESS (even when bulk ESS looks fine) signals that the sampler is stuck in the bulk and rarely reaches the extremes of the posterior.


``` python
import altair as alt

diag_data = [
    {"coefficient": i, "r_hat": float(mr.r_hat[i]), "ess": float(mr.ess[i])}
    for i in range(len(mr.r_hat))
]

threshold = alt.Chart({"values": [{}]}).mark_rule(
    strokeDash=[4, 4], color="firebrick"
).encode(y=alt.datum(1.1))

points = alt.Chart({"values": diag_data}).mark_point(size=60, filled=True).encode(
    x=alt.X("coefficient:O", title="Coefficient index"),
    y=alt.Y("r_hat:Q", title="R-hat", scale=alt.Scale(zero=False)),
    color=alt.condition(
        alt.datum.r_hat > 1.1,
        alt.value("firebrick"),
        alt.value("steelblue"),
    ),
)

(points + threshold).properties(
    title="R-hat by coefficient (red = above 1.1 threshold)",
    width=500,
    height=250,
)
```


<style>
  #altair-viz-f5064bb67a67400b921f4e26fa172411.vega-embed {
    width: 100%;
    display: flex;
  }

  #altair-viz-f5064bb67a67400b921f4e26fa172411.vega-embed details,
  #altair-viz-f5064bb67a67400b921f4e26fa172411.vega-embed details summary {
    position: relative;
  }
</style>


## Divergences

A **divergent transition** occurs when the leapfrog integrator encounters a region of very high curvature and the Hamiltonian energy error exceeds 1000. This signals that the sampler has strayed into a part of the posterior that the step size cannot handle correctly. The resulting draw is biased (and not just noisy).

`n_divergent` counts the total number of divergent transitions across all chains and all post-warmup samples. For a well-specified model with a reasonable step size, this should be zero. Any non-zero value is a warning that posterior geometry is causing problems:

``` python
if mr.n_divergent > 0:
    print(f"Warning: {mr.n_divergent} divergent transitions detected.")
    print("Try a higher target_accept (e.g. 0.90) or reparameterize the model.")
```

Increasing `target_accept` causes the dual-averaging algorithm to adapt to a smaller step size, which reduces energy errors at the cost of shorter trajectories. If divergences persist despite a high target acceptance rate, the posterior may have a funnel-shaped geometry that requires reparameterization.


# Predictions and uncertainty

`predict(se=True)` and confidence intervals work exactly as with any other fitting method. Standard errors are derived from the posterior samples:


``` python
x_new = np.linspace(0, 2 * np.pi, 200)
preds = model.predict({"x": x_new}, se=True, interval="confidence")

obs_data = [{"x": float(x[i]), "y": float(y[i])} for i in range(n)]
points = alt.Chart({"values": obs_data}).mark_circle(
    size=15, opacity=0.3, color="steelblue"
).encode(
    x=alt.X("x:Q", title="x"),
    y=alt.Y("y:Q", title="count"),
)

fit_data = [
    {
        "x": float(x_new[i]),
        "fit": float(preds.values[i]),
        "lower": float(preds.lower[i]),
        "upper": float(preds.upper[i]),
    }
    for i in range(len(x_new))
]

line = alt.Chart({"values": fit_data}).mark_line(
    color="firebrick", strokeWidth=2
).encode(x="x:Q", y="fit:Q")

band = alt.Chart({"values": fit_data}).mark_area(
    opacity=0.2, color="firebrick"
).encode(x="x:Q", y="lower:Q", y2="upper:Q")

(band + line + points).properties(
    title="Poisson GAM with MCMC: fitted curve and 95% credible interval",
    width=500,
    height=300,
)
```


<style>
  #altair-viz-87a6759c7653497484bcf19ee4b626cf.vega-embed {
    width: 100%;
    display: flex;
  }

  #altair-viz-87a6759c7653497484bcf19ee4b626cf.vega-embed details,
  #altair-viz-87a6759c7653497484bcf19ee4b626cf.vega-embed details summary {
    position: relative;
  }
</style>


# Posterior samples

`posterior_samples(n)` draws coefficient vectors from the empirical posterior:


``` python
beta_samples = model.posterior_samples(n=500, seed=23)
print(f"Shape: {beta_samples.shape}")  # (p, 500)
```


    Shape: (10, 500)


This also works after `method="REML"`, `"GCV"`, or `"VI"` fits (in those cases samples are drawn from the Laplace or variational posterior), so code that calls [posterior_samples()](../reference/GAM.md#whittaker.GAM.posterior_samples) is inference-method agnostic.


# Simulating from the posterior

`simulate()` propagates posterior coefficient uncertainty through the response distribution, giving a posterior predictive sample:


``` python
sims = model.simulate(n_sim=200, seed=0)
print(f"Simulations shape: {sims.shape}")  # (n, 200): integer count draws
```


    Simulations shape: (300, 200)


# Controlling MCMC

Pass an `mcmc_options=` dict to `fit()` to override any sampler setting:

``` python
model = wk.GAM("y ~ s(x)", family=Poisson()).fit(
    data,
    method="MCMC",
    mcmc_options={
        "n_chains": 4,          # number of independent chains (default: 4)
        "n_samples": 1000,      # post-warmup draws per chain (default: 1000)
        "n_warmup": 500,        # warmup (adaptation) draws per chain (default: 500)
        "max_tree_depth": 10,   # NUTS: max binary-tree doublings per step (default: 10)
        "seed": 23,             # seed for reproducibility
    },
)
```


## Choosing a sampler

Two samplers are available via the `sampler` key:

- **`"NUTS"`** (default): No-U-Turn Sampler. Automatically selects trajectory length by growing a binary tree until a U-turn is detected. Requires no manual tuning of leapfrog steps and generally mixes better than fixed-length HMC.
- **`"HMC"`**: Static-trajectory HMC with a fixed number of leapfrog steps per proposal, controlled by `leapfrog_steps` (the default is `10`).

``` python
# Static HMC with 20 leapfrog steps per proposal
model = wk.GAM("y ~ s(x)", family=Poisson()).fit(
    data,
    method="MCMC",
    mcmc_options={"sampler": "HMC", "leapfrog_steps": 20},
)
```

`mean_tree_depth` in `MCMCResult` reports the mean number of binary-tree doublings per post-warmup NUTS step (0.0 for HMC). Each doubling doubles the number of leapfrog evaluations: depth j corresponds to 2^j steps. A depth of 5-7 is typical; if it consistently hits `max_tree_depth` consider raising that limit.


## Step-size adaptation

Whittaker uses dual-averaging step-size adaptation (Nesterov 2009, as implemented in Stan) during the warmup phase. The adapted step size is fixed for the sampling phase. The default target acceptance rate is `0.65`. NUTS often benefits from a higher target:

``` python
model = wk.GAM("y ~ s(x)", family=Poisson()).fit(
    data,
    method="MCMC",
    mcmc_options={"target_accept": 0.80},  # default: 0.65
)
```

For NUTS, `acceptance_rate` in `MCMCResult` reports the mean per-leaf acceptance statistic (mean \min(1, \exp(H_0 - H_i)) over all leapfrog steps), which tracks the dual-averaging target and is directly comparable to HMC's Metropolis acceptance rate.

Higher target acceptance rates lead to smaller step sizes and more correlated draws. Lower rates lead to larger steps but more rejections.


## Mass matrix

Whittaker uses a **two-phase warmup** to adapt the diagonal mass matrix.

**Phase 1** (first half of warmup): the sampler runs with a mass matrix initialized from the Laplace posterior covariance, M = \text{diag}(1 / V\_\beta), collecting draws while the dual-averaging algorithm adapts the step size.

**Midpoint**: the empirical variance of the phase-1 draws is used to update the mass matrix, M \leftarrow \text{diag}(1 / \widehat{\sigma}^2\_\beta). The step-size dual-averaging is reset so that phase 2 can re-adapt the step size to the new geometry. This update is skipped when fewer than 10 warmup steps precede it, so very short warmup schedules are handled gracefully.

**Phase 2** (second half of warmup): dual-averaging continues with the updated mass matrix, and the averaged step size from the end of this phase is used for all post-warmup sampling.

This two-phase scheme substantially improves mixing when posterior coefficient scales differ by orders of magnitude. And this is common with B-spline bases where an intercept at scale e^5 coexists with smooth components at scale 0.01.


# Properties not available for MCMC fits

The `deviance`, [null_deviance](../reference/GAM.md#whittaker.GAM.null_deviance), [deviance_explained](../reference/GAM.md#whittaker.GAM.deviance_explained), `aic`, `bic`, and [gcv_score](../reference/GAM.md#whittaker.GAM.gcv_score) properties are not defined for MCMC fits and raise `NotImplementedError`. Use R-hat and ESS for convergence assessment, and posterior predictive checks (via `simulate()`) for model comparison.


# Where to go next

- **[Model comparison with LOO](loo.md)**: use PSIS-LOO to compare the predictive accuracy of competing Bayesian GAM fits.
- **[Posterior predictive distributions](posterior-predict.md)**: the full predictive distribution at new data points, including observation noise.
- **[Posterior predictive checks](ppc.md)**: assess whether a model generates data consistent with the observations.
- **[Variational inference](variational-inference.md)**: a faster alternative to MCMC for non-Gaussian families that gives better-calibrated intervals than the Laplace approximation.
- **[Prediction and inference](prediction.md)**: confidence intervals, simultaneous bands, and term-level predictions from any fitted model.
- **[Model diagnostics](diagnostics.md)**: residual plots and `model.check()`.
- **[Response families](families.md)**: how the family affects posterior shape and MCMC efficiency.
