PyMC Samplers

PyMC Samplers Skill

You are an expert in PyMC sampling methods, helping migrate notebook code to work with the current stable PyMC version.

Core Sampling Functions

pm.sample() - Main Sampling Function

The primary function for MCMC sampling:

import pymc as pm

with pm.Model() as model:
    # Define model...

    # Basic sampling
    trace = pm.sample(
        draws=1000,           # Number of samples per chain
        tune=1000,            # Number of tuning steps
        chains=4,             # Number of chains
        cores=4,              # Parallel chains
        random_seed=None  # Set for reproducibility        # Reproducibility
    )

Full pm.sample() Parameters

trace = pm.sample(
    draws=1000,                    # Samples to draw
    tune=1000,                     # Tuning/burn-in steps
    chains=4,                      # Number of chains
    cores=None,                    # CPU cores (None = all available)
    random_seed=None,              # Random seed for reproducibility
    step=None,                     # Custom step sampler(s)
    initvals=None,                 # Initial values
    init="auto",                   # Initialization: "auto", "jitter+adapt_diag", etc.
    n_init=200000,                 # Samples for initialization
    progressbar=True,              # Show progress bar
    return_inferencedata=True,     # Return ArviZ InferenceData
    idata_kwargs=None,             # InferenceData creation kwargs
    compute_convergence_checks=True, # Compute Rhat, ESS
    discard_tuned_samples=True,    # Exclude tuning samples
    target_accept=0.8,             # For NUTS: target acceptance rate
    callback=None,                 # Callback function per draw
)

Sampling Step Methods

NUTS (No-U-Turn Sampler)

Default and recommended for most models with continuous variables:

# Automatic (default)
trace = pm.sample(1000)

# Explicit NUTS configuration
trace = pm.sample(
    1000,
    target_accept=0.95,  # Higher = slower but more accurate
    max_treedepth=10     # Maximum tree depth
)

# Manual NUTS setup
from pymc.step_methods import NUTS

with model:
    step = NUTS()
    trace = pm.sample(1000, step=step)

Metropolis Samplers

For simpler models or discrete variables:

from pymc.step_methods import Metropolis

with model:
    # Basic Metropolis
    step = Metropolis()
    trace = pm.sample(5000, step=step)

    # With specific proposal distribution
    step = Metropolis(proposal_dist=pm.NormalProposal)

BinaryMetropolis

Optimized for binary variables:

from pymc.step_methods import BinaryMetropolis

with model:
    x = pm.Bernoulli("x", p=0.5)
    step = BinaryMetropolis([x])
    trace = pm.sample(5000, step=step)

Slice Sampler

For unimodal distributions:

from pymc.step_methods import Slice

with model:
    step = Slice()
    trace = pm.sample(5000, step=step)

HamiltonianMC

Standard Hamiltonian Monte Carlo:

from pymc.step_methods import HamiltonianMC

with model:
    step = HamiltonianMC()
    trace = pm.sample(1000, step=step)

CompoundStep - Multiple Samplers

Combine different samplers for different variables:

from pymc.step_methods import NUTS, BinaryMetropolis, CompoundStep

with model:
    continuous_vars = [mu, sigma]
    discrete_vars = [z]

    # NUTS for continuous, Metropolis for discrete
    step1 = NUTS(vars=continuous_vars)
    step2 = BinaryMetropolis(vars=discrete_vars)

    # Compound step
    step = CompoundStep([step1, step2])
    trace = pm.sample(1000, step=step)

Prior and Posterior Predictive Sampling

Prior Predictive Checks

Sample from the prior before seeing data:

with model:
    # Sample from prior
    prior_predictive = pm.sample_prior_predictive(
        samples=500,
        random_seed=None  # Set for reproducibility
    )

# Inspect prior predictions
import arviz as az
az.plot_ppc(prior_predictive, group="prior")

Posterior Predictive Checks

Sample predictions after fitting:

with model:
    # Fit model
    trace = pm.sample(1000)

    # Generate posterior predictions
    posterior_predictive = pm.sample_posterior_predictive(
        trace,
        var_names=None,        # Variables to sample (None = all)
        random_seed=None  # Set for reproducibility,
        progressbar=True
    )

# Check model fit
az.plot_ppc(posterior_predictive)

Out-of-Sample Predictions

# Update data for new predictions
with model:
    pm.set_data({"X": X_test})

    # Sample predictions
    predictions = pm.sample_posterior_predictive(
        trace,
        var_names=["y"]  # Only prediction variable
    )

Initialization Methods

Initialization Strategies

# Auto (default) - usually "jitter+adapt_diag"
trace = pm.sample(1000, init="auto")

# Jitter around prior mean + adapt diagonal mass matrix
trace = pm.sample(1000, init="jitter+adapt_diag")

# Jitter around prior mean + full adaptation
trace = pm.sample(1000, init="jitter+adapt_full")

# ADVI initialization
trace = pm.sample(1000, init="adapt_diag")

# Maximum a posteriori (MAP) initialization
trace = pm.sample(1000, init="adapt_diag_grad")

Custom Initial Values

# Specify starting values
initvals = {
    "mu": 0.5,
    "sigma": 1.0
}

trace = pm.sample(1000, initvals=initvals)

NUTS Initialization Helper

# Initialize NUTS sampler with specific settings
from pymc import init_nuts

with model:
    init_vals, step = init_nuts(
        init="auto",
        chains=4,
        random_seed=None  # Set for reproducibility,
        progressbar=True
    )

    trace = pm.sample(1000, step=step, initvals=init_vals)

Drawing Single Samples

# Draw single random sample from variable
with model:
    sample = pm.draw(mu, draws=1)

    # Multiple draws
    samples = pm.draw(mu, draws=100)

    # Draw multiple variables
    sample_dict = pm.draw([mu, sigma], draws=100)

Computing Deterministics

After sampling, compute deterministic quantities:

with model:
    trace = pm.sample(1000)

    # Compute deterministic variables from trace
    trace = pm.compute_deterministics(trace)

JAX-Based Samplers

BlackJAX NUTS

from pymc.sampling.jax import sample_blackjax_nuts

with model:
    trace = sample_blackjax_nuts(
        draws=1000,
        tune=1000,
        chains=4,
        target_accept=0.8
    )

NumPyro NUTS

from pymc.sampling.jax import sample_numpyro_nuts

with model:
    trace = sample_numpyro_nuts(
        draws=1000,
        tune=1000,
        chains=4,
        target_accept=0.8
    )

Common Migration Issues

PyMC3 → latest PyMC version

1. Return type changes

# Old (PyMC3) - returns MultiTrace
trace = pm.sample(1000)

# New (latest PyMC version) - returns InferenceData by default
idata = pm.sample(1000, return_inferencedata=True)  # default

# Access posterior samples
idata.posterior

# Get old-style trace if needed
trace = pm.sample(1000, return_inferencedata=False)

2. Step method imports

# Old
from pymc3.step_methods import NUTS, Metropolis

# New
from pymc.step_methods import NUTS, Metropolis

3. Tuning parameter

# Old: some versions used 'n_tune'
trace = pm.sample(draws=1000, n_tune=500)

# New: use 'tune'
trace = pm.sample(draws=1000, tune=500)

4. Initialization

# Old: testval in variable definition
x = pm.Normal("x", mu=0, sigma=1, testval=0.5)

# New: use initvals in sample()
x = pm.Normal("x", mu=0, sigma=1)
trace = pm.sample(1000, initvals={"x": 0.5})

Sampling Best Practices

1. Use NUTS for continuous models - It's efficient and robust
2. Run multiple chains - At least 4 chains to check convergence
3. Check convergence - Inspect Rhat (< 1.1) and ESS (> 200)
4. Tune adequately - Default 1000 steps usually sufficient
5. Adjust target_accept if needed - Increase if divergences occur
6. Use random_seed - For reproducibility
7. Prior predictive checks - Always check priors make sense
8. Posterior predictive checks - Validate model fit

Troubleshooting Sampling Issues

Divergences

# Increase target acceptance rate
trace = pm.sample(1000, target_accept=0.95)

# Increase max tree depth
trace = pm.sample(1000, max_treedepth=15)

# Reparameterize model (non-centered → centered or vice versa)

Slow Sampling

# Use fewer chains on more cores
trace = pm.sample(1000, chains=2, cores=2)

# Try JAX backend
trace = sample_numpyro_nuts(1000)

# Reduce tuning steps if model is simple
trace = pm.sample(1000, tune=500)

Poor Mixing

# Increase tuning period
trace = pm.sample(1000, tune=2000)

# Try different initialization
trace = pm.sample(1000, init="adapt_diag")

# Reparameterize the model
# (e.g., non-centered parameterization for hierarchical models)

High Rhat or Low ESS

# Run longer chains
trace = pm.sample(5000, tune=2000)

# More chains
trace = pm.sample(1000, chains=8)

# Check model specification
# - Are priors appropriate?
# - Is model identified?
# - Are there correlations that need reparameterization?

Example Usage

import pymc as pm
import numpy as np
import arviz as az

# Generate data
np.random.seed(42)
true_mu = 3.0
true_sigma = 1.5
data = np.random.normal(true_mu, true_sigma, size=100)

# Build model
with pm.Model() as model:
    # Priors
    mu = pm.Normal("mu", mu=0, sigma=10)
    sigma = pm.HalfNormal("sigma", sigma=5)

    # Likelihood
    y = pm.Normal("y", mu=mu, sigma=sigma, observed=data)

    # Prior predictive check
    prior_pred = pm.sample_prior_predictive(samples=500)

    # Sample posterior
    trace = pm.sample(
        draws=1000,
        tune=1000,
        chains=4,
        cores=4,
        target_accept=0.9,
        random_seed=None  # Set for reproducibility
    )

    # Posterior predictive check
    post_pred = pm.sample_posterior_predictive(trace)

# Check convergence
print(az.summary(trace, var_names=["mu", "sigma"]))

# Visualize
az.plot_trace(trace, var_names=["mu", "sigma"])
az.plot_posterior(trace, var_names=["mu", "sigma"])
az.plot_ppc(post_pred)

When to Use This Skill

• Configuring sampling parameters
• Fixing divergences or sampling issues
• Choosing appropriate step methods
• Converting PyMC3 sampling code
• Implementing prior/posterior predictive checks
• Troubleshooting convergence problems
• Optimizing sampling performance
• Working with mixed discrete/continuous models