quivers.diagnostics.arviz_io

MCMCResult → xarray.DataTree conversion.

arviz_io

Conversion from quivers fits to xarray.DataTree (the ArviZ 1.x container).

The MCMC sampler hands back quivers.inference.MCMCResult records whose samples dict already follows the (num_chains, num_samples, *site_shape) shape convention ArviZ expects. This module repackages those tensors as numpy arrays under the canonical ArviZ group names (posterior, sample_stats, posterior_predictive, log_likelihood, observed_data) plus the user-supplied coords / dims for nice axis labels in plots.

to_datatree

to_datatree(posterior: MCMCResult, *, observed_data: Mapping[str, Tensor] | None = None, posterior_predictive: Mapping[str, Tensor] | None = None, log_likelihood: Mapping[str, Tensor] | None = None, constant_data: Mapping[str, Tensor] | None = None, coords: Mapping[str, list] | None = None, dims: Mapping[str, list[str]] | None = None) -> DataTree

Convert an MCMCResult into an ArviZ-style DataTree.

PARAMETER	DESCRIPTION
posterior	Sampler output. posterior.samples[name] carries an (num_chains, num_samples, *site_shape) tensor; each becomes a posterior variable. log_densities populates sample_stats/lp; acceptance_rates and divergence_counts populate sample_stats. TYPE: MCMCResult
observed_data	Site name to observed tensor (the original data used at fit time). Becomes the observed_data group. TYPE: Mapping[str, Tensor] DEFAULT: None
posterior_predictive	Site name to posterior-predictive draws of shape (num_chains, num_samples, *site_shape). Becomes the posterior_predictive group. TYPE: Mapping[str, Tensor] DEFAULT: None
log_likelihood	Site name to per-observation log-likelihood of shape (num_chains, num_samples, *obs_shape). Becomes the log_likelihood group; required for arviz.loo / arviz.waic. TYPE: Mapping[str, Tensor] DEFAULT: None
constant_data	Site name to fixed covariate tensor (e.g. design matrix). Becomes the constant_data group. TYPE: Mapping[str, Tensor] DEFAULT: None
coords	Coordinate values per named axis (e.g. {"Verb": ["eat", "drink", "run"]}). Forwarded to ArviZ. TYPE: Mapping[str, list] DEFAULT: None
dims	Per-site axis names (e.g. {"beta": ["Verb"]}). Forwarded to ArviZ. TYPE: Mapping[str, list[str]] DEFAULT: None

RETURNS	DESCRIPTION
DataTree	Canonical ArviZ DataTree consumable by every plotting and diagnostic function in the ArviZ 1.x API.

Source code in src/quivers/diagnostics/arviz_io.py

def to_datatree(
    posterior: MCMCResult,
    *,
    observed_data: Mapping[str, torch.Tensor] | None = None,
    posterior_predictive: Mapping[str, torch.Tensor] | None = None,
    log_likelihood: Mapping[str, torch.Tensor] | None = None,
    constant_data: Mapping[str, torch.Tensor] | None = None,
    coords: Mapping[str, list] | None = None,
    dims: Mapping[str, list[str]] | None = None,
) -> xr.DataTree:
    """Convert an `MCMCResult` into an ArviZ-style DataTree.

    Parameters
    ----------
    posterior : MCMCResult
        Sampler output.  ``posterior.samples[name]`` carries an
        ``(num_chains, num_samples, *site_shape)`` tensor; each
        becomes a posterior variable.  ``log_densities`` populates
        ``sample_stats/lp``; ``acceptance_rates`` and
        ``divergence_counts`` populate ``sample_stats``.
    observed_data : Mapping[str, torch.Tensor], optional
        Site name to observed tensor (the original data used at
        fit time).  Becomes the ``observed_data`` group.
    posterior_predictive : Mapping[str, torch.Tensor], optional
        Site name to posterior-predictive draws of shape
        ``(num_chains, num_samples, *site_shape)``.  Becomes the
        ``posterior_predictive`` group.
    log_likelihood : Mapping[str, torch.Tensor], optional
        Site name to per-observation log-likelihood of shape
        ``(num_chains, num_samples, *obs_shape)``.  Becomes the
        ``log_likelihood`` group; required for
        `arviz.loo` / `arviz.waic`.
    constant_data : Mapping[str, torch.Tensor], optional
        Site name to fixed covariate tensor (e.g. design matrix).
        Becomes the ``constant_data`` group.
    coords : Mapping[str, list], optional
        Coordinate values per named axis (e.g.
        ``{"Verb": ["eat", "drink", "run"]}``).  Forwarded to ArviZ.
    dims : Mapping[str, list[str]], optional
        Per-site axis names (e.g. ``{"beta": ["Verb"]}``).
        Forwarded to ArviZ.

    Returns
    -------
    xr.DataTree
        Canonical ArviZ DataTree consumable by every plotting and
        diagnostic function in the ArviZ 1.x API.
    """
    data: dict[str, dict] = {}

    posterior_group: dict = {}
    for name, t in posterior.samples.items():
        posterior_group[name] = _tensor_to_numpy(t)
    data["posterior"] = posterior_group

    sample_stats_group: dict = {
        "lp": _tensor_to_numpy(posterior.log_densities),
    }
    # acceptance_rate and diverging are per-chain scalars in
    # quivers; ArviZ expects per-draw arrays.  We broadcast.
    n_chains = posterior.num_chains
    n_samples = posterior.num_samples
    sample_stats_group["mean_acceptance_per_chain"] = (
        posterior.acceptance_rates.detach()
        .cpu()
        .numpy()
        .reshape(n_chains, 1)
        .repeat(n_samples, axis=1)
    )
    sample_stats_group["total_divergences_per_chain"] = (
        posterior.divergence_counts.detach()
        .cpu()
        .numpy()
        .reshape(n_chains, 1)
        .repeat(n_samples, axis=1)
    )
    data["sample_stats"] = sample_stats_group

    if observed_data:
        data["observed_data"] = {
            name: _tensor_to_numpy(t) for name, t in observed_data.items()
        }
    if posterior_predictive:
        data["posterior_predictive"] = {
            name: _tensor_to_numpy(t) for name, t in posterior_predictive.items()
        }
    if log_likelihood:
        data["log_likelihood"] = {
            name: _tensor_to_numpy(t) for name, t in log_likelihood.items()
        }
    if constant_data:
        data["constant_data"] = {
            name: _tensor_to_numpy(t) for name, t in constant_data.items()
        }

    return az.from_dict(
        data,
        coords=dict(coords) if coords is not None else None,
        dims=dict(dims) if dims is not None else None,
    )

to_datatree_from_svi

to_datatree_from_svi(samples: Mapping[str, Tensor], *, log_densities: Tensor | None = None, observed_data: Mapping[str, Tensor] | None = None, posterior_predictive: Mapping[str, Tensor] | None = None, log_likelihood: Mapping[str, Tensor] | None = None, constant_data: Mapping[str, Tensor] | None = None, coords: Mapping[str, list] | None = None, dims: Mapping[str, list[str]] | None = None) -> DataTree

Convert a set of SVI or Predictive draws into an ArviZ-style DataTree with a single pseudo-chain.

Each entry in samples is a (num_samples, *site_shape) tensor (the standard shape produced by Predictive(model, guide, num_samples=...)). A leading chain dimension of size 1 is inserted so the resulting DataTree matches the (chain, draw, *event) shape ArviZ expects. log_densities, when supplied, becomes sample_stats/lp; ArviZ diagnostic functions (loo, waic, plot_trace) work uniformly on the resulting object.

Source code in src/quivers/diagnostics/arviz_io.py

def to_datatree_from_svi(
    samples: Mapping[str, torch.Tensor],
    *,
    log_densities: torch.Tensor | None = None,
    observed_data: Mapping[str, torch.Tensor] | None = None,
    posterior_predictive: Mapping[str, torch.Tensor] | None = None,
    log_likelihood: Mapping[str, torch.Tensor] | None = None,
    constant_data: Mapping[str, torch.Tensor] | None = None,
    coords: Mapping[str, list] | None = None,
    dims: Mapping[str, list[str]] | None = None,
) -> xr.DataTree:
    """Convert a set of SVI or Predictive draws into an ArviZ-style
    DataTree with a single pseudo-chain.

    Each entry in `samples` is a `(num_samples, *site_shape)` tensor
    (the standard shape produced by `Predictive(model, guide,
    num_samples=...)`). A leading chain dimension of size 1 is
    inserted so the resulting DataTree matches the `(chain, draw,
    *event)` shape ArviZ expects. `log_densities`, when supplied,
    becomes `sample_stats/lp`; ArviZ diagnostic functions (`loo`,
    `waic`, `plot_trace`) work uniformly on the resulting object.
    """
    data: dict[str, dict] = {}

    posterior_group: dict = {}
    for name, t in samples.items():
        arr = _tensor_to_numpy(t)
        if arr.ndim == 0:
            arr = arr.reshape(1)
        posterior_group[name] = arr[None, ...]
    data["posterior"] = posterior_group

    if log_densities is not None:
        lp = _tensor_to_numpy(log_densities)
        if lp.ndim == 0:
            lp = lp.reshape(1)
        data["sample_stats"] = {"lp": lp[None, ...]}

    if observed_data:
        data["observed_data"] = {
            name: _tensor_to_numpy(t) for name, t in observed_data.items()
        }
    if posterior_predictive:
        data["posterior_predictive"] = {
            name: _tensor_to_numpy(t)[None, ...]
            for name, t in posterior_predictive.items()
        }
    if log_likelihood:
        data["log_likelihood"] = {
            name: _tensor_to_numpy(t)[None, ...] for name, t in log_likelihood.items()
        }
    if constant_data:
        data["constant_data"] = {
            name: _tensor_to_numpy(t) for name, t in constant_data.items()
        }

    return az.from_dict(
        data,
        coords=dict(coords) if coords is not None else None,
        dims=dict(dims) if dims is not None else None,
    )

to_datatree_any

to_datatree_any(fit, *, observed_data: Mapping[str, Tensor] | None = None, posterior_predictive: Mapping[str, Tensor] | None = None, log_likelihood: Mapping[str, Tensor] | None = None, constant_data: Mapping[str, Tensor] | None = None, coords: Mapping[str, list] | None = None, dims: Mapping[str, list[str]] | None = None) -> DataTree

Dispatch a DataTree conversion on the fit type.

Accepts:

• MCMCResult — routes to the original to_datatree chain / draw semantics.
• A plain dict of str -> Tensor (SVI / Predictive draws) — routes to to_datatree_from_svi with a single pseudo-chain.
• A tuple (samples, log_densities) — same as the dict form but populates sample_stats/lp from the second entry.

Consumers with a custom fit-container type can add a case here or call the specific to_datatree / to_datatree_from_svi entry point directly. The generic dispatch closes issue #43 by giving SVI users a one-line ArviZ export path.

Source code in src/quivers/diagnostics/arviz_io.py

def to_datatree_any(
    fit,
    *,
    observed_data: Mapping[str, torch.Tensor] | None = None,
    posterior_predictive: Mapping[str, torch.Tensor] | None = None,
    log_likelihood: Mapping[str, torch.Tensor] | None = None,
    constant_data: Mapping[str, torch.Tensor] | None = None,
    coords: Mapping[str, list] | None = None,
    dims: Mapping[str, list[str]] | None = None,
) -> xr.DataTree:
    """Dispatch a DataTree conversion on the fit type.

    Accepts:

    * [`MCMCResult`][quivers.inference.MCMCResult] — routes to the
      original `to_datatree` chain / draw semantics.
    * A plain dict of `str -> Tensor` (SVI / Predictive draws) —
      routes to `to_datatree_from_svi` with a single pseudo-chain.
    * A tuple `(samples, log_densities)` — same as the dict form
      but populates `sample_stats/lp` from the second entry.

    Consumers with a custom fit-container type can add a case here
    or call the specific `to_datatree` / `to_datatree_from_svi`
    entry point directly. The generic dispatch closes issue #43 by
    giving SVI users a one-line ArviZ export path.
    """
    if isinstance(fit, MCMCResult):
        return to_datatree(
            fit,
            observed_data=observed_data,
            posterior_predictive=posterior_predictive,
            log_likelihood=log_likelihood,
            constant_data=constant_data,
            coords=coords,
            dims=dims,
        )
    if isinstance(fit, tuple) and len(fit) == 2:
        samples, log_densities = fit
        return to_datatree_from_svi(
            samples,
            log_densities=log_densities,
            observed_data=observed_data,
            posterior_predictive=posterior_predictive,
            log_likelihood=log_likelihood,
            constant_data=constant_data,
            coords=coords,
            dims=dims,
        )
    if isinstance(fit, Mapping):
        return to_datatree_from_svi(
            fit,
            observed_data=observed_data,
            posterior_predictive=posterior_predictive,
            log_likelihood=log_likelihood,
            constant_data=constant_data,
            coords=coords,
            dims=dims,
        )
    raise TypeError(
        f"to_datatree_any: cannot dispatch on fit type "
        f"{type(fit).__name__}; supported: MCMCResult, "
        "Mapping[str, Tensor] (SVI / Predictive), "
        "tuple[Mapping, Tensor] (samples + log-densities)."
    )