"""BayesianResult dataclass and related types for Bayesian inference output.
This module contains the result container returned by BayesianMixin.fit_bayesian(),
including ArviZ InferenceData conversion with fast-path xarray assembly.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, TypedDict
import numpy as np
from rheojax.logging import get_logger
logger = get_logger(__name__)
if TYPE_CHECKING:
from numpyro.infer import MCMC
class DiagnosticsDict(TypedDict, total=False):
"""Typed structure for Bayesian convergence diagnostics."""
r_hat: dict[str, float]
ess: dict[str, float]
divergences: int
diagnostics_valid: bool
total_samples: int
num_chains: int
num_samples_per_chain: int
error: str
init_strategy: str
warm_start_failed: bool
bfmi: float
nonfinite_draws: int
[docs]
@dataclass(eq=False)
class BayesianResult:
"""Results from Bayesian inference with NUTS sampling.
This dataclass stores the complete results of NumPyro NUTS sampling,
including posterior samples, summary statistics, convergence diagnostics,
and placeholders for future model comparison metrics.
Attributes:
posterior_samples: Dictionary mapping parameter names to arrays of
posterior samples (shape: [num_samples * num_chains, ]). All arrays are float64.
summary: Dictionary with summary statistics for each parameter.
Contains nested dicts with 'mean', 'std', and quantiles.
diagnostics: Dictionary with convergence diagnostics including:
- r_hat: Gelman-Rubin statistic for each parameter (dict)
- ess: Effective sample size for each parameter (dict)
- divergences: Number of divergent transitions (int)
num_samples: Number of posterior samples per chain (after warmup).
num_chains: Number of MCMC chains used in sampling.
mcmc: NumPyro MCMC object containing full sampling information including
NUTS-specific diagnostics (energy, divergences, tree depth).
Required for ArviZ visualization with full diagnostics.
model_comparison: Dictionary for model comparison metrics (WAIC, LOO).
Currently a placeholder for future implementation.
_inference_data: Cached ArviZ InferenceData object. Automatically
created on first access via to_inference_data(). Do not set manually.
Example:
>>> result = model.fit_bayesian(X, y)
>>> print(result.summary["a"]["mean"])
>>> print(result.diagnostics["r_hat"]["a"])
>>> # Convert to ArviZ InferenceData for advanced plotting
>>> idata = result.to_inference_data()
"""
posterior_samples: dict[str, np.ndarray]
summary: dict[str, dict[str, float]]
diagnostics: DiagnosticsDict
num_samples: int
num_chains: int
mcmc: MCMC | None = None
model_comparison: dict[str, float] = field(default_factory=dict)
_inference_data: Any | None = field(default=None, repr=False)
_inference_data_ll: Any | None = field(default=None, repr=False)
[docs]
def __post_init__(self):
"""Validate result after initialization."""
logger.debug(
"Initializing BayesianResult",
num_parameters=len(self.posterior_samples),
num_samples=self.num_samples,
num_chains=self.num_chains,
)
# Ensure posterior_samples are float64 numpy arrays
# BAY-05: skip copy when already float64 NumPy to avoid eager allocation
for name, samples in self.posterior_samples.items():
arr = np.asarray(samples)
if arr.dtype != np.float64:
arr = arr.astype(np.float64)
self.posterior_samples[name] = arr
logger.debug(
"BayesianResult initialized",
parameter_names=list(self.posterior_samples.keys()),
)
[docs]
def __eq__(self, other: object) -> bool:
"""Compare results by value, using np.array_equal for posterior_samples.
The default dataclass __eq__ would crash (ValueError) whenever any
posterior_samples array has more than one element, because tuple/dict
equality forces bool() on the ndarray comparison result. mcmc and the
cached InferenceData objects are excluded as not meaningfully
comparable via value equality.
"""
if not isinstance(other, BayesianResult):
return NotImplemented
if self.posterior_samples.keys() != other.posterior_samples.keys():
return False
if not all(
np.array_equal(self.posterior_samples[k], other.posterior_samples[k])
for k in self.posterior_samples
):
return False
return (
self.summary == other.summary
and self.diagnostics == other.diagnostics
and self.num_samples == other.num_samples
and self.num_chains == other.num_chains
and self.model_comparison == other.model_comparison
)
[docs]
def to_inference_data(self, log_likelihood: bool = False) -> Any:
"""Convert to ArviZ InferenceData format for advanced visualization.
Converts the NumPyro MCMC result to ArviZ InferenceData format, which
enables access to ArviZ's comprehensive plotting and diagnostic tools.
The conversion preserves all NUTS-specific diagnostics including energy,
divergences, and tree depth information.
The InferenceData object is cached after first conversion to avoid
repeated conversion overhead. The ``log_likelihood=False`` and
``log_likelihood=True`` variants are cached independently.
Args:
log_likelihood: If True, compute pointwise log-likelihood for
LOO model comparison (az.loo()). ArviZ 1.x has no WAIC
implementation. This re-evaluates the model for all samples
(~600-800ms slower). Default False for faster conversion
when only plotting.
Returns:
ArviZ InferenceData object containing:
- posterior: Posterior samples for all parameters
- sample_stats: NUTS diagnostics (energy, divergences, etc.)
- log_likelihood: Only when ``log_likelihood=True``
- Additional groups as available from NumPyro
Raises:
ImportError: If arviz is not installed
ValueError: If MCMC object was not stored (older results)
Example:
>>> result = model.fit_bayesian(X, y)
>>> idata = result.to_inference_data() # Fast: no log-lik
>>> az.plot_trace(idata)
>>>
>>> # For model comparison (slower):
>>> idata_ll = result.to_inference_data(log_likelihood=True)
>>> az.loo(idata_ll)
Note:
Requires arviz package: pip install arviz
The MCMC object must be present (automatically stored by fit_bayesian).
"""
logger.debug(
"Converting BayesianResult to InferenceData",
log_likelihood=log_likelihood,
)
# Return cached version if available
if log_likelihood and self._inference_data_ll is not None:
logger.debug("Returning cached InferenceData (with log_likelihood)")
return self._inference_data_ll
if not log_likelihood and self._inference_data is not None:
logger.debug("Returning cached InferenceData (without log_likelihood)")
return self._inference_data
# Ensure MCMC object is available
if self.mcmc is None:
logger.error("MCMC object not available for InferenceData conversion")
raise ValueError(
"MCMC object not available for conversion. "
"This may be a result from an older version. "
"Re-run fit_bayesian() to generate a compatible result."
)
# Import arviz (lazy import)
from rheojax.core.arviz_utils import import_arviz
try:
import_arviz()
logger.debug("ArviZ imported successfully")
except ImportError as exc:
logger.error("ArviZ import failed", exc_info=True)
raise ImportError(
"ArviZ is required for InferenceData conversion. "
"Install it with: pip install arviz"
) from exc
if log_likelihood:
# log_likelihood=True requires model re-evaluation (slow path).
# Delegate to az.from_numpyro which traces the model to extract
# pointwise log-likelihoods for WAIC/LOO computation.
logger.debug(
"Creating InferenceData from MCMC object (with log_likelihood)",
)
try:
az_full = import_arviz(required=("from_numpyro",))
idata = az_full.from_numpyro(self.mcmc, log_likelihood=True)
except ImportError as exc:
raise ImportError(
"ArviZ is required for log-likelihood computation. "
"Install it with: pip install arviz"
) from exc
# ArviZ 1.x's PSIS-LOO backend (arviz_stats) mutates array
# buffers in place; az.from_numpyro() leaves log_likelihood as
# a live JAX array, which rejects in-place mutation (TypeError:
# JAX arrays are immutable). Materialize to plain numpy so
# az.loo()/az.compare()/az.plot_khat() work on this idata.
for var in idata.log_likelihood.data_vars:
idata.log_likelihood[var].values = np.asarray(
idata.log_likelihood[var].values
)
logger.info(
"InferenceData created successfully (with log_likelihood)",
num_chains=self.num_chains,
num_samples=self.num_samples,
)
self._inference_data_ll = idata
return idata
# Fast path (log_likelihood=False): build InferenceData directly from
# numpy arrays already present in BayesianResult, bypassing the
# az.from_numpyro model-trace that triggers XLA recompilation
# (~500-1500ms on first call). This reduces conversion to <5ms.
#
# ArviZ's from_numpyro calls numpyro.handlers.trace().get_trace() in
# NumPyroConverter.__init__ to discover observed sites for the
# observed_data / log_likelihood groups. When log_likelihood=False we
# don't need those groups, so we can skip the trace entirely and
# assemble the two groups we do need (posterior, sample_stats) from
# data that is already on the host.
logger.debug("Building InferenceData directly (fast path, no model trace)")
import xarray as xr
# --- posterior group ---
# Use get_samples(group_by_chain=True) to include ALL sampled sites
# (model parameters + deterministic sites like num_nonfinite) so the
# result matches what az.from_numpyro would return.
# self.posterior_samples only contains param_names + sigma params;
# deterministic sites (numpyro.deterministic) are excluded from it.
num_chains = self.num_chains
posterior_dict: dict[str, xr.DataArray] = {}
try:
# Prefer group_by_chain samples — already (num_chains, num_draws)
_raw_samples = self.mcmc.get_samples(group_by_chain=True)
if hasattr(_raw_samples, "_asdict"):
_raw_samples = _raw_samples._asdict()
for name, arr in _raw_samples.items():
np_arr = np.asarray(arr)
if np_arr.ndim >= 2:
# Shape is already (num_chains, num_draws[, ...])
posterior_dict[name] = xr.DataArray(
np_arr,
dims=("chain", "draw")
+ tuple(f"dim_{i}" for i in range(np_arr.ndim - 2)),
)
else:
# Fallback: reshape flat array
try:
shaped = np_arr.reshape(num_chains, -1)
except ValueError:
shaped = np_arr[np.newaxis, :]
posterior_dict[name] = xr.DataArray(shaped, dims=("chain", "draw"))
except Exception as exc:
# Fall back to posterior_samples (already numpy, already on host).
# Reset any partially-populated entries so a mid-loop failure
# can't leave an internally-inconsistent subset of fields
# silently attached, and log at WARNING (not DEBUG) so the drop
# is visible at normal verbosity instead of only in debug logs
# (mirrors the sample_stats fallback below).
posterior_dict = {}
logger.warning(
"get_samples(group_by_chain=True) failed, using posterior_samples",
error=str(exc),
)
for name, flat_arr in self.posterior_samples.items():
try:
shaped = flat_arr.reshape(num_chains, -1)
except ValueError:
shaped = flat_arr[np.newaxis, :]
posterior_dict[name] = xr.DataArray(shaped, dims=("chain", "draw"))
# --- sample_stats group ---
# Mirrors ArviZ's NumPyroConverter.sample_stats_to_xarray() rename map.
_stat_rename = {
"potential_energy": "lp",
"adapt_state.step_size": "step_size",
"num_steps": "n_steps",
"accept_prob": "acceptance_rate",
}
stats_dict: dict[str, xr.DataArray] = {}
try:
try:
extra_fields: dict[str, Any] = self.mcmc.get_extra_fields(
group_by_chain=True
)
except TypeError:
extra_fields = self.mcmc.get_extra_fields()
if isinstance(extra_fields, dict):
for stat, value in extra_fields.items():
if isinstance(value, (dict, tuple)):
logger.debug(
"Skipping sample_stats field: unsupported container type",
stat=stat,
value_type=type(value).__name__,
)
continue
arr = np.asarray(value)
# Ensure (chain, draw) shape
if arr.ndim == 1:
try:
arr = arr.reshape(num_chains, -1)
except ValueError:
arr = arr[np.newaxis, :]
elif arr.ndim != 2:
logger.debug(
"Skipping sample_stats field: unexpected ndim",
stat=stat,
ndim=arr.ndim,
)
continue # Skip unexpected shapes
# Match ArviZ's real NumPyroConverter.sample_stats_to_xarray():
# potential_energy -> "lp" verbatim, no sign flip. "energy" (the
# true Hamiltonian = potential + kinetic) is passed through as-is
# when NumPyro provided it via extra_fields; see fallback below
# for older MCMC objects that only requested potential_energy.
if stat == "potential_energy":
stats_dict["lp"] = xr.DataArray(arr, dims=("chain", "draw"))
elif stat == "energy":
stats_dict["energy"] = xr.DataArray(arr, dims=("chain", "draw"))
else:
dest_name = _stat_rename.get(stat, stat)
stats_dict[dest_name] = xr.DataArray(
arr, dims=("chain", "draw")
)
# Fallback only: older MCMC objects that never requested the real
# "energy" extra field. Approximate it from potential_energy alone
# and say so — this proxy omits kinetic energy, so BFMI computed
# from it is approximate, not the exact E-BFMI statistic.
if "energy" not in stats_dict and "lp" in stats_dict:
logger.warning(
"NumPyro extra_fields did not include 'energy' (true "
"Hamiltonian); approximating sample_stats['energy'] from "
"potential_energy alone. BFMI computed from this proxy "
"will be approximate.",
)
stats_dict["energy"] = xr.DataArray(
stats_dict["lp"].values, dims=("chain", "draw")
)
except Exception as exc:
# sample_stats is a best-effort group here (posterior-only callers,
# e.g. WAIC/LOO consumers, must not break on an unrelated failure
# in this block). Reset any partially-populated stats so a
# mid-loop failure can't leave an internally-inconsistent subset
# of fields silently attached, and log at WARNING (not DEBUG) so
# the drop is visible at normal verbosity instead of only in
# debug logs.
stats_dict = {}
logger.warning(
"Failed to extract sample_stats from MCMC extra fields; "
"sample_stats group will be absent from InferenceData",
error=str(exc),
)
posterior_np = {name: da.values for name, da in posterior_dict.items()}
stats_np = {name: da.values for name, da in stats_dict.items()}
from rheojax.core.arviz_utils import inference_data_from_dict
idata = inference_data_from_dict(
{
"posterior": posterior_np if posterior_np else None,
"sample_stats": stats_np if stats_np else None,
}
)
logger.info(
"InferenceData created successfully (fast path)",
num_chains=self.num_chains,
num_samples=self.num_samples,
)
self._inference_data = idata
return idata
__all__ = [
"BayesianResult",
"DiagnosticsDict",
]