"""Sampler specifications for posterior inference."""
import os
from typing import Literal
import arviz as az
import pymc as pm
from pydantic import Field
from impulso._base import ImpulsoModel
def _default_nuts_sampler() -> Literal["pymc", "nutpie"]:
"""Return 'nutpie' if installed, otherwise 'pymc'."""
try:
import nutpie # noqa: F401
except ImportError:
return "pymc"
else:
return "nutpie"
def _default_progressbar() -> bool:
"""Show the sampler progress bar, except during documentation builds.
Sphinx sets ``IMPULSO_DOCS_BUILD=1`` so rendered notebooks do not embed the
live progress widget. Normal usage is unaffected.
"""
return os.environ.get("IMPULSO_DOCS_BUILD") != "1"
[docs]
class NUTSSampler(ImpulsoModel):
"""NUTS sampler configuration for PyMC.
Attributes:
draws: Number of posterior draws per chain.
tune: Number of tuning steps per chain.
chains: Number of independent chains.
cores: Number of CPU cores. None = auto-detect.
target_accept: Target acceptance rate for NUTS.
random_seed: Random seed for reproducibility.
nuts_sampler: NUTS backend. Auto-detects nutpie if installed.
progressbar: Show the sampler progress bar. Defaults to True, but off
during documentation builds (``IMPULSO_DOCS_BUILD=1``).
nuts_sampler_kwargs: Extra keyword arguments forwarded verbatim to
the NUTS backend (`pm.sample(nuts_sampler_kwargs=...)`). Useful
for backend-specific adaptation options — e.g. nutpie's
`low_rank_modified_mass_matrix=True`, which handles the
ill-conditioned posteriors that arise in large VARs with many
near-collinear lag regressors, where diagonal mass-matrix
adaptation mixes poorly.
"""
draws: int = Field(1000, ge=1)
tune: int = Field(1000, ge=0)
chains: int = Field(4, ge=1)
cores: int | None = Field(None, ge=1)
target_accept: float = Field(0.8, gt=0, lt=1)
random_seed: int | None = None
nuts_sampler: Literal["pymc", "nutpie"] = Field(default_factory=_default_nuts_sampler)
progressbar: bool = Field(default_factory=_default_progressbar)
nuts_sampler_kwargs: dict | None = None
[docs]
def sample(self, model: pm.Model) -> az.InferenceData:
"""Run NUTS sampling on the given PyMC model.
Args:
model: A fully specified PyMC model.
Returns:
ArviZ InferenceData with posterior and log_likelihood groups.
"""
with model:
idata = pm.sample(
draws=self.draws,
tune=self.tune,
chains=self.chains,
cores=self.cores,
target_accept=self.target_accept,
random_seed=self.random_seed,
nuts_sampler=self.nuts_sampler,
progressbar=self.progressbar,
nuts_sampler_kwargs=self.nuts_sampler_kwargs or {},
idata_kwargs={"log_likelihood": True},
)
return idata