Source code for

"""emcee-specific conversion code."""

import warnings
from collections import OrderedDict

import numpy as np
import xarray as xr

from .. import utils
from .base import dict_to_dataset, generate_dims_coords, make_attrs
from .inference_data import InferenceData

def _verify_names(sampler, var_names, arg_names, slices):
    """Make sure var_names and arg_names are assigned reasonably.

    This is meant to run before loading emcee objects into InferenceData.
    In case var_names or arg_names is None, will provide defaults. If they are
    not None, it verifies there are the right number of them.

    Throws a ValueError in case validation fails.

    sampler : emcee.EnsembleSampler
        Fitted emcee sampler
    var_names : list[str] or None
        Names for the emcee parameters
    arg_names : list[str] or None
        Names for the args/observations provided to emcee
    slices : list[seq] or None
        slices to select the variables (used for multidimensional variables)

    list[str], list[str], list[seq]
        Defaults for var_names, arg_names and slices
    # There are 3 possible cases: emcee2, emcee3 and sampler read from h5 file (emcee3 only)
    if hasattr(sampler, "args"):
        ndim = sampler.chain.shape[-1]
        num_args = len(sampler.args)
    elif hasattr(sampler, "log_prob_fn"):
        ndim = sampler.get_chain().shape[-1]
        num_args = len(sampler.log_prob_fn.args)
        ndim = sampler.get_chain().shape[-1]
        num_args = 0  # emcee only stores the posterior samples

    if slices is None:
        slices = utils.arange(ndim)
        num_vars = ndim
        num_vars = len(slices)
    indices = utils.arange(ndim)
    slicing_try = np.concatenate([utils.one_de(indices[idx]) for idx in slices])
    if len(set(slicing_try)) != ndim:
            "Check slices: Not all parameters in chain captured. "
            f"{ndim} are present, and {len(slicing_try)} have been captured.",
    if len(slicing_try) != len(set(slicing_try)):
        warnings.warn(f"Overlapping slices. Check the index present: {slicing_try}", UserWarning)

    if var_names is None:
        var_names = [f"var_{idx}" for idx in range(num_vars)]
    if arg_names is None:
        arg_names = [f"arg_{idx}" for idx in range(num_args)]

    if len(var_names) != num_vars:
        raise ValueError(
            f"The sampler has {num_vars} variables, "
            f"but only {len(var_names)} var_names were provided!"

    if len(arg_names) != num_args:
        raise ValueError(
            f"The sampler has {num_args} args, "
            f"but only {len(arg_names)} arg_names were provided!"
    return var_names, arg_names, slices

# pylint: disable=too-many-instance-attributes
class EmceeConverter:
    """Encapsulate emcee specific logic."""

    def __init__(
        var_names, arg_names, slices = _verify_names(sampler, var_names, arg_names, slices)
        self.sampler = sampler
        self.var_names = var_names
        self.slices = slices
        self.arg_names = arg_names
        self.arg_groups = arg_groups
        self.blob_names = blob_names
        self.blob_groups = blob_groups
        self.index_origin = index_origin
        self.coords = coords
        self.dims = dims
        import emcee

        self.emcee = emcee

    def posterior_to_xarray(self):
        """Convert the posterior to an xarray dataset."""
        # Use emcee3 syntax, else use emcee2
        if hasattr(self.sampler, "get_chain"):
            samples_ary = self.sampler.get_chain().swapaxes(0, 1)
            samples_ary = self.sampler.chain

        data = {
            var_name: (samples_ary[(..., idx)])
            for idx, var_name in zip(self.slices, self.var_names)
        return dict_to_dataset(

    def args_to_xarray(self):
        """Convert emcee args to observed and constant_data xarray Datasets."""
        dims = {} if self.dims is None else self.dims
        if self.arg_groups is None:
            self.arg_groups = ["observed_data" for _ in self.arg_names]
        if len(self.arg_names) != len(self.arg_groups):
            raise ValueError(
                "arg_names and arg_groups must have the same length, or arg_groups be None"
        arg_groups_set = set(self.arg_groups)
        bad_groups = [
            group for group in arg_groups_set if group not in ("observed_data", "constant_data")
        if bad_groups:
            raise SyntaxError(
                "all arg_groups values should be either 'observed_data' or 'constant_data' , "
                f"not {bad_groups}"
        obs_const_dict = {group: OrderedDict() for group in arg_groups_set}
        for idx, (arg_name, group) in enumerate(zip(self.arg_names, self.arg_groups)):
            # Use emcee3 syntax, else use emcee2
            arg_array = np.atleast_1d(
                if hasattr(self.sampler, "log_prob_fn")
                else self.sampler.args[idx]
            arg_dims = dims.get(arg_name)
            arg_dims, coords = generate_dims_coords(
            # filter coords based on the dims
            coords = {key: xr.IndexVariable((key,), data=coords[key]) for key in arg_dims}
            obs_const_dict[group][arg_name] = xr.DataArray(arg_array, dims=arg_dims, coords=coords)
        for key, values in obs_const_dict.items():
            obs_const_dict[key] = xr.Dataset(data_vars=values, attrs=make_attrs(library=self.emcee))
        return obs_const_dict

    def blobs_to_dict(self):
        """Convert blobs to dictionary {groupname: xr.Dataset}.

        It also stores lp values in sample_stats group.
        store_blobs = self.blob_names is not None
        self.blob_names = [] if self.blob_names is None else self.blob_names
        if self.blob_groups is None:
            self.blob_groups = ["log_likelihood" for _ in self.blob_names]
        if len(self.blob_names) != len(self.blob_groups):
            raise ValueError(
                "blob_names and blob_groups must have the same length, or blob_groups be None"
        if store_blobs:
            if int(self.emcee.__version__[0]) >= 3:
                blobs = self.sampler.get_blobs()
                blobs = np.array(self.sampler.blobs, dtype=object)
            if (blobs is None or blobs.size == 0) and self.blob_names:
                raise ValueError("No blobs in sampler, blob_names must be None")
            if len(blobs.shape) == 2:
                blobs = np.expand_dims(blobs, axis=-1)
            blobs = blobs.swapaxes(0, 2)
            nblobs, nwalkers, ndraws, *_ = blobs.shape
            if len(self.blob_names) != nblobs and len(self.blob_names) > 1:
                raise ValueError(
                    "Incorrect number of blob names. "
                    f"Expected {nblobs}, found {len(self.blob_names)}"
        blob_groups_set = set(self.blob_groups)
        idata_groups = ("posterior", "observed_data", "constant_data")
        if np.any(np.isin(list(blob_groups_set), idata_groups)):
            raise SyntaxError(
                f"{idata_groups} groups should not come from blobs. "
                "Using them here would overwrite their actual values"
        blob_dict = {group: OrderedDict() for group in blob_groups_set}
        if len(self.blob_names) == 1:
            blob_dict[self.blob_groups[0]][self.blob_names[0]] = blobs.swapaxes(0, 2).swapaxes(0, 1)
            for i_blob, (name, group) in enumerate(zip(self.blob_names, self.blob_groups)):
                # for coherent blobs (all having the same dimensions) one line is enough
                blob = blobs[i_blob]
                # for blobs of different size, we get an array of arrays, which we convert
                # to an ndarray per blob_name
                if blob.dtype == object:
                    blob = blob.reshape(-1)
                    blob = np.stack(blob)
                    blob = blob.reshape((nwalkers, ndraws, -1))
                blob_dict[group][name] = np.squeeze(blob)

        # store lp in sample_stats group
        blob_dict["sample_stats"]["lp"] = (
            self.sampler.get_log_prob().swapaxes(0, 1)
            if hasattr(self.sampler, "get_log_prob")
            else self.sampler.lnprobability
        for key, values in blob_dict.items():
            blob_dict[key] = dict_to_dataset(
        return blob_dict

    def to_inference_data(self):
        """Convert all available data to an InferenceData object."""
        blobs_dict = self.blobs_to_dict()
        obs_const_dict = self.args_to_xarray()
        return InferenceData(
            **{"posterior": self.posterior_to_xarray(), **obs_const_dict, **blobs_dict}

[docs] def from_emcee( sampler=None, var_names=None, slices=None, arg_names=None, arg_groups=None, blob_names=None, blob_groups=None, index_origin=None, coords=None, dims=None, ): """Convert emcee data into an InferenceData object. For a usage example read :ref:`emcee_conversion` Parameters ---------- sampler : emcee.EnsembleSampler Fitted sampler from emcee. var_names : list of str, optional A list of names for variables in the sampler slices : list of array-like or slice, optional A list containing the indexes of each variable. Should only be used for multidimensional variables. arg_names : list of str, optional A list of names for args in the sampler arg_groups : list of str, optional A list of the group names (either ``observed_data`` or ``constant_data``) where args in the sampler are stored. If None, all args will be stored in observed data group. blob_names : list of str, optional A list of names for blobs in the sampler. When None, blobs are omitted, independently of them being present in the sampler or not. blob_groups : list of str, optional A list of the groups where blob_names variables should be assigned respectively. If blob_names!=None and blob_groups is None, all variables are assigned to log_likelihood group coords : dict of {str : array_like}, optional Map of dimensions to coordinates dims : dict of {str : list of str}, optional Map variable names to their coordinates Returns ------- arviz.InferenceData """ return EmceeConverter( sampler=sampler, var_names=var_names, slices=slices, arg_names=arg_names, arg_groups=arg_groups, blob_names=blob_names, blob_groups=blob_groups, index_origin=index_origin, coords=coords, dims=dims, ).to_inference_data()