arviz 0.23.3__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

arviz/data/inference_data.py

@@ -1,2386 +0,0 @@
-# pylint: disable=too-many-lines,too-many-public-methods
-"""Data structure for using netcdf groups with xarray."""
-import os
-import re
-import sys
-import uuid
-import warnings
-from collections import OrderedDict, defaultdict
-from collections.abc import MutableMapping, Sequence
-from copy import copy as ccopy
-from copy import deepcopy
-import datetime
-from html import escape
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Iterable,
-    Iterator,
-    List,
-    Mapping,
-    Optional,
-    Tuple,
-    TypeVar,
-    Union,
-    overload,
-)
-
-import numpy as np
-import xarray as xr
-from packaging import version
-
-from ..rcparams import rcParams
-from ..utils import HtmlTemplate, _subset_list, _var_names, either_dict_or_kwargs
-from .base import _extend_xr_method, _make_json_serializable, dict_to_dataset
-
-if sys.version_info[:2] >= (3, 9):
-    # As of 3.9, collections.abc types support generic parameters themselves.
-    from collections.abc import ItemsView, ValuesView
-else:
-    # These typing imports are deprecated in 3.9, and moved to collections.abc instead.
-    from typing import ItemsView, ValuesView
-
-if TYPE_CHECKING:
-    from typing_extensions import Literal
-
-try:
-    import ujson as json
-except ImportError:
-    # mypy struggles with conditional imports expressed as catching ImportError:
-    # https://github.com/python/mypy/issues/1153
-    import json  # type: ignore
-
-
-SUPPORTED_GROUPS = [
-    "posterior",
-    "posterior_predictive",
-    "predictions",
-    "log_likelihood",
-    "log_prior",
-    "sample_stats",
-    "prior",
-    "prior_predictive",
-    "sample_stats_prior",
-    "observed_data",
-    "constant_data",
-    "predictions_constant_data",
-    "unconstrained_posterior",
-    "unconstrained_prior",
-]
-
-WARMUP_TAG = "warmup_"
-
-SUPPORTED_GROUPS_WARMUP = [
-    f"{WARMUP_TAG}posterior",
-    f"{WARMUP_TAG}posterior_predictive",
-    f"{WARMUP_TAG}predictions",
-    f"{WARMUP_TAG}sample_stats",
-    f"{WARMUP_TAG}log_likelihood",
-    f"{WARMUP_TAG}log_prior",
-]
-
-SUPPORTED_GROUPS_ALL = SUPPORTED_GROUPS + SUPPORTED_GROUPS_WARMUP
-
-InferenceDataT = TypeVar("InferenceDataT", bound="InferenceData")
-
-
-def _compressible_dtype(dtype):
-    """Check basic dtypes for automatic compression."""
-    if dtype.kind == "V":
-        return all(_compressible_dtype(item) for item, _ in dtype.fields.values())
-    return dtype.kind in {"b", "i", "u", "f", "c", "S"}
-
-
-class InferenceData(Mapping[str, xr.Dataset]):
-    """Container for inference data storage using xarray.
-
-    For a detailed introduction to ``InferenceData`` objects and their usage, see
-    :ref:`xarray_for_arviz`. This page provides help and documentation
-    on ``InferenceData`` methods and their low level implementation.
-    """
-
-    def __init__(
-        self,
-        attrs: Union[None, Mapping[Any, Any]] = None,
-        warn_on_custom_groups: bool = False,
-        **kwargs: Union[xr.Dataset, List[xr.Dataset], Tuple[xr.Dataset, xr.Dataset]],
-    ) -> None:
-        """Initialize InferenceData object from keyword xarray datasets.
-
-        Parameters
-        ----------
-        attrs : dict
-            sets global attribute for InferenceData object.
-        warn_on_custom_groups : bool, default False
-            Emit a warning when custom groups are present in the InferenceData.
-            "custom group" means any group whose name isn't defined in :ref:`schema`
-        kwargs :
-            Keyword arguments of xarray datasets
-
-        Examples
-        --------
-        Initiate an InferenceData object from scratch, not recommended. InferenceData
-        objects should be initialized using ``from_xyz`` methods, see :ref:`data_api` for more
-        details.
-
-        .. ipython::
-
-            In [1]: import arviz as az
-               ...: import numpy as np
-               ...: import xarray as xr
-               ...: dataset = xr.Dataset(
-               ...:     {
-               ...:         "a": (["chain", "draw", "a_dim"], np.random.normal(size=(4, 100, 3))),
-               ...:         "b": (["chain", "draw"], np.random.normal(size=(4, 100))),
-               ...:     },
-               ...:     coords={
-               ...:         "chain": (["chain"], np.arange(4)),
-               ...:         "draw": (["draw"], np.arange(100)),
-               ...:         "a_dim": (["a_dim"], ["x", "y", "z"]),
-               ...:     }
-               ...: )
-               ...: idata = az.InferenceData(posterior=dataset, prior=dataset)
-               ...: idata
-
-        We have created an ``InferenceData`` object with two groups. Now we can check its
-        contents:
-
-        .. ipython::
-
-            In [1]: idata.posterior
-
-        """
-        self._groups: List[str] = []
-        self._groups_warmup: List[str] = []
-        self._attrs: Union[None, dict] = dict(attrs) if attrs is not None else None
-        save_warmup = kwargs.pop("save_warmup", False)
-        key_list = [key for key in SUPPORTED_GROUPS_ALL if key in kwargs]
-        for key in kwargs:
-            if key not in SUPPORTED_GROUPS_ALL:
-                key_list.append(key)
-                if warn_on_custom_groups:
-                    warnings.warn(
-                        f"{key} group is not defined in the InferenceData scheme", UserWarning
-                    )
-        for key in key_list:
-            dataset = kwargs[key]
-            dataset_warmup = None
-            if dataset is None:
-                continue
-            elif isinstance(dataset, (list, tuple)):
-                dataset, dataset_warmup = dataset
-            elif not isinstance(dataset, xr.Dataset):
-                raise ValueError(
-                    "Arguments to InferenceData must be xarray Datasets "
-                    f"(argument '{key}' was type '{type(dataset)}')"
-                )
-            if not key.startswith(WARMUP_TAG):
-                if dataset:
-                    setattr(self, key, dataset)
-                    self._groups.append(key)
-            elif key.startswith(WARMUP_TAG):
-                if dataset:
-                    setattr(self, key, dataset)
-                    self._groups_warmup.append(key)
-            if save_warmup and dataset_warmup is not None and dataset_warmup:
-                key = f"{WARMUP_TAG}{key}"
-                setattr(self, key, dataset_warmup)
-                self._groups_warmup.append(key)
-
-    @property
-    def attrs(self) -> dict:
-        """Attributes of InferenceData object."""
-        if self._attrs is None:
-            self._attrs = {}
-        return self._attrs
-
-    @attrs.setter
-    def attrs(self, value) -> None:
-        self._attrs = dict(value)
-
-    def __repr__(self) -> str:
-        """Make string representation of InferenceData object."""
-        msg = "Inference data with groups:\n\t> {options}".format(
-            options="\n\t> ".join(self._groups)
-        )
-        if self._groups_warmup:
-            msg += f"\n\nWarmup iterations saved ({WARMUP_TAG}*)."
-        return msg
-
-    def _repr_html_(self) -> str:
-        """Make html representation of InferenceData object."""
-        try:
-            from xarray.core.options import OPTIONS
-
-            display_style = OPTIONS["display_style"]
-            if display_style == "text":
-                html_repr = f"<pre>{escape(repr(self))}</pre>"
-            else:
-                elements = "".join(
-                    [
-                        HtmlTemplate.element_template.format(
-                            group_id=group + str(uuid.uuid4()),
-                            group=group,
-                            xr_data=getattr(  # pylint: disable=protected-access
-                                self, group
-                            )._repr_html_(),
-                        )
-                        for group in self._groups_all
-                    ]
-                )
-                formatted_html_template = (  # pylint: disable=possibly-unused-variable
-                    HtmlTemplate.html_template.format(elements)
-                )
-                css_template = HtmlTemplate.css_template  # pylint: disable=possibly-unused-variable
-                html_repr = f"{locals()['formatted_html_template']}{locals()['css_template']}"
-        except:  # pylint: disable=bare-except
-            html_repr = f"<pre>{escape(repr(self))}</pre>"
-        return html_repr
-
-    def __delattr__(self, group: str) -> None:
-        """Delete a group from the InferenceData object."""
-        if group in self._groups:
-            self._groups.remove(group)
-        elif group in self._groups_warmup:
-            self._groups_warmup.remove(group)
-        object.__delattr__(self, group)
-
-    def __delitem__(self, key: str) -> None:
-        """Delete an item from the InferenceData object using del idata[key]."""
-        self.__delattr__(key)
-
-    @property
-    def _groups_all(self) -> List[str]:
-        return self._groups + self._groups_warmup
-
-    def __len__(self) -> int:
-        """Return the number of groups in this InferenceData object."""
-        return len(self._groups_all)
-
-    def __iter__(self) -> Iterator[str]:
-        """Iterate over groups in InferenceData object."""
-        yield from self._groups_all
-
-    def __contains__(self, key: object) -> bool:
-        """Return True if the named item is present, and False otherwise."""
-        return key in self._groups_all
-
-    def __getitem__(self, key: str) -> xr.Dataset:
-        """Get item by key."""
-        if key not in self._groups_all:
-            raise KeyError(key)
-        return getattr(self, key)
-
-    def __setitem__(self, key: str, value: xr.Dataset):
-        """Set item by key and update group list accordingly."""
-        if key.startswith(WARMUP_TAG):
-            self._groups_warmup.append(key)
-        else:
-            self._groups.append(key)
-        setattr(self, key, value)
-
-    def groups(self) -> List[str]:
-        """Return all groups present in InferenceData object."""
-        return self._groups_all
-
-    class InferenceDataValuesView(ValuesView[xr.Dataset]):
-        """ValuesView implementation for InferenceData, to allow it to implement Mapping."""
-
-        def __init__(  # pylint: disable=super-init-not-called
-            self, parent: "InferenceData"
-        ) -> None:
-            """Create a new InferenceDataValuesView from an InferenceData object."""
-            self.parent = parent
-
-        def __len__(self) -> int:
-            """Return the number of groups in the parent InferenceData."""
-            return len(self.parent._groups_all)
-
-        def __iter__(self) -> Iterator[xr.Dataset]:
-            """Iterate through the Xarray datasets present in the InferenceData object."""
-            parent = self.parent
-            for group in parent._groups_all:
-                yield getattr(parent, group)
-
-        def __contains__(self, key: object) -> bool:
-            """Return True if the given Xarray dataset is one of the values, and False otherwise."""
-            if not isinstance(key, xr.Dataset):
-                return False
-
-            for dataset in self:
-                if dataset.equals(key):
-                    return True
-
-            return False
-
-    def values(self) -> "InferenceData.InferenceDataValuesView":
-        """Return a view over the Xarray Datasets present in the InferenceData object."""
-        return InferenceData.InferenceDataValuesView(self)
-
-    class InferenceDataItemsView(ItemsView[str, xr.Dataset]):
-        """ItemsView implementation for InferenceData, to allow it to implement Mapping."""
-
-        def __init__(  # pylint: disable=super-init-not-called
-            self, parent: "InferenceData"
-        ) -> None:
-            """Create a new InferenceDataItemsView from an InferenceData object."""
-            self.parent = parent
-
-        def __len__(self) -> int:
-            """Return the number of groups in the parent InferenceData."""
-            return len(self.parent._groups_all)
-
-        def __iter__(self) -> Iterator[Tuple[str, xr.Dataset]]:
-            """Iterate through the groups and corresponding Xarray datasets in the InferenceData."""
-            parent = self.parent
-            for group in parent._groups_all:
-                yield group, getattr(parent, group)
-
-        def __contains__(self, key: object) -> bool:
-            """Return True if the (group, dataset) tuple is present, and False otherwise."""
-            parent = self.parent
-            if not isinstance(key, tuple) or len(key) != 2:
-                return False
-
-            group, dataset = key
-            if group not in parent._groups_all:
-                return False
-
-            if not isinstance(dataset, xr.Dataset):
-                return False
-
-            existing_dataset = getattr(parent, group)
-            return existing_dataset.equals(dataset)
-
-    def items(self) -> "InferenceData.InferenceDataItemsView":
-        """Return a view over the groups and datasets present in the InferenceData object."""
-        return InferenceData.InferenceDataItemsView(self)
-
-    @staticmethod
-    def from_netcdf(
-        filename,
-        *,
-        engine="h5netcdf",
-        group_kwargs=None,
-        regex=False,
-        base_group: str = "/",
-    ) -> "InferenceData":
-        """Initialize object from a netcdf file.
-
-        Expects that the file will have groups, each of which can be loaded by xarray.
-        By default, the datasets of the InferenceData object will be lazily loaded instead
-        of being loaded into memory. This
-        behaviour is regulated by the value of ``az.rcParams["data.load"]``.
-
-        Parameters
-        ----------
-        filename : str
-            location of netcdf file
-        engine : {"h5netcdf", "netcdf4"}, default "h5netcdf"
-            Library used to read the netcdf file.
-        group_kwargs : dict of {str: dict}, optional
-            Keyword arguments to be passed into each call of :func:`xarray.open_dataset`.
-            The keys of the higher level should be group names or regex matching group
-            names, the inner dicts re passed to ``open_dataset``
-            This feature is currently experimental.
-        regex : bool, default False
-            Specifies where regex search should be used to extend the keyword arguments.
-            This feature is currently experimental.
-        base_group : str, default "/"
-            The group in the netCDF file where the InferenceData is stored. By default,
-            assumes that the file only contains an InferenceData object.
-
-        Returns
-        -------
-        InferenceData
-        """
-        groups = {}
-        attrs = {}
-
-        if engine == "h5netcdf":
-            import h5netcdf
-        elif engine == "netcdf4":
-            import netCDF4 as nc
-        else:
-            raise ValueError(
-                f"Invalid value for engine: {engine}. Valid options are: h5netcdf or netcdf4"
-            )
-
-        try:
-            with (
-                h5netcdf.File(filename, mode="r")
-                if engine == "h5netcdf"
-                else nc.Dataset(filename, mode="r")
-            ) as file_handle:
-                if base_group == "/":
-                    data = file_handle
-                else:
-                    data = file_handle[base_group]
-
-                data_groups = list(data.groups)
-
-            for group in data_groups:
-                group_kws = {}
-
-                group_kws = {}
-                if group_kwargs is not None and regex is False:
-                    group_kws = group_kwargs.get(group, {})
-                if group_kwargs is not None and regex is True:
-                    for key, kws in group_kwargs.items():
-                        if re.search(key, group):
-                            group_kws = kws
-                group_kws.setdefault("engine", engine)
-                data = xr.open_dataset(filename, group=f"{base_group}/{group}", **group_kws)
-                if rcParams["data.load"] == "eager":
-                    with data:
-                        groups[group] = data.load()
-                else:
-                    groups[group] = data
-
-            with xr.open_dataset(filename, engine=engine, group=base_group) as data:
-                attrs.update(data.load().attrs)
-
-            return InferenceData(attrs=attrs, **groups)
-        except OSError as err:
-            if err.errno == -101:
-                raise type(err)(
-                    str(err)
-                    + (
-                        " while reading a NetCDF file. This is probably an error in HDF5, "
-                        "which happens because your OS does not support HDF5 file locking.  See "
-                        "https://stackoverflow.com/questions/49317927/"
-                        "errno-101-netcdf-hdf-error-when-opening-netcdf-file#49317928"
-                        " for a possible solution."
-                    )
-                ) from err
-            raise err
-
-    def to_netcdf(
-        self,
-        filename: str,
-        compress: bool = True,
-        groups: Optional[List[str]] = None,
-        engine: str = "h5netcdf",
-        base_group: str = "/",
-        overwrite_existing: bool = True,
-    ) -> str:
-        """Write InferenceData to netcdf4 file.
-
-        Parameters
-        ----------
-        filename : str
-            Location to write to
-        compress : bool, optional
-            Whether to compress result. Note this saves disk space, but may make
-            saving and loading somewhat slower (default: True).
-        groups : list, optional
-            Write only these groups to netcdf file.
-        engine : {"h5netcdf", "netcdf4"}, default "h5netcdf"
-            Library used to read the netcdf file.
-        base_group : str, default "/"
-            The group in the netCDF file where the InferenceData is will be stored.
-            By default, will write to the root of the netCDF file
-        overwrite_existing : bool, default True
-            Whether to overwrite the existing file or append to it.
-
-        Returns
-        -------
-        str
-            Location of netcdf file
-        """
-        if base_group is None:
-            base_group = "/"
-
-        if os.path.exists(filename) and not overwrite_existing:
-            mode = "a"
-        else:
-            mode = "w"  # overwrite first, then append
-
-        if self._attrs:
-            xr.Dataset(attrs=self._attrs).to_netcdf(
-                filename, mode=mode, engine=engine, group=base_group
-            )
-            mode = "a"
-
-        if self._groups_all:  # check's whether a group is present or not.
-            if groups is None:
-                groups = self._groups_all
-            else:
-                groups = [group for group in self._groups_all if group in groups]
-
-            for group in groups:
-                data = getattr(self, group)
-                kwargs = {"engine": engine}
-                if compress:
-                    kwargs["encoding"] = {
-                        var_name: {"zlib": True}
-                        for var_name, values in data.variables.items()
-                        if _compressible_dtype(values.dtype)
-                    }
-                data.to_netcdf(filename, mode=mode, group=f"{base_group}/{group}", **kwargs)
-                data.close()
-                mode = "a"
-        elif not self._attrs:  # creates a netcdf file for an empty InferenceData object.
-            if engine == "h5netcdf":
-                import h5netcdf
-
-                empty_netcdf_file = h5netcdf.File(filename, mode="w")
-            elif engine == "netcdf4":
-                import netCDF4 as nc
-
-                empty_netcdf_file = nc.Dataset(filename, mode="w", format="NETCDF4")
-            empty_netcdf_file.close()
-        return filename
-
-    def to_datatree(self):
-        """Convert InferenceData object to a :class:`~xarray.DataTree`."""
-        try:
-            from xarray import DataTree
-        except ImportError as err:
-            raise ImportError(
-                "xarray must be have DataTree in order to use InferenceData.to_datatree. "
-                "Update to xarray>=2024.11.0"
-            ) from err
-        dt = DataTree.from_dict({group: ds for group, ds in self.items()})
-        dt.attrs = self.attrs
-        return dt
-
-    @staticmethod
-    def from_datatree(datatree):
-        """Create an InferenceData object from a :class:`~xarray.DataTree`.
-
-        Parameters
-        ----------
-        datatree : DataTree
-        """
-        return InferenceData(
-            attrs=datatree.attrs,
-            **{group: child.to_dataset() for group, child in datatree.children.items()},
-        )
-
-    def to_dict(self, groups=None, filter_groups=None):
-        """Convert InferenceData to a dictionary following xarray naming conventions.
-
-        Parameters
-        ----------
-        groups : list, optional
-            Groups where the transformation is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups: {None, "like", "regex"}, optional, default=None
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-
-        Returns
-        -------
-        dict
-            A dictionary containing all groups of InferenceData object.
-            When `data=False` return just the schema.
-        """
-        ret = defaultdict(dict)
-        if self._groups_all:  # check's whether a group is present or not.
-            if groups is None:
-                groups = self._group_names(groups, filter_groups)
-            else:
-                groups = [group for group in self._groups_all if group in groups]
-
-            for group in groups:
-                dataset = getattr(self, group)
-                data = {}
-                for var_name, dataarray in dataset.items():
-                    data[var_name] = dataarray.values
-                    dims = []
-                    for coord_name, coord_values in dataarray.coords.items():
-                        if coord_name not in ("chain", "draw") and not coord_name.startswith(
-                            f"{var_name}_dim_"
-                        ):
-                            dims.append(coord_name)
-                            ret["coords"][coord_name] = coord_values.values
-
-                    if group in (
-                        "predictions",
-                        "predictions_constant_data",
-                    ):
-                        dims_key = "pred_dims"
-                    else:
-                        dims_key = "dims"
-                    if len(dims) > 0:
-                        ret[dims_key][var_name] = dims
-                    ret[group] = data
-                ret[f"{group}_attrs"] = dataset.attrs
-
-        ret["attrs"] = self.attrs
-        return ret
-
-    def to_json(self, filename, groups=None, filter_groups=None, **kwargs):
-        """Write InferenceData to a json file.
-
-        Parameters
-        ----------
-        filename : str
-            Location to write to
-        groups : list, optional
-            Groups where the transformation is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups: {None, "like", "regex"}, optional, default=None
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        kwargs : dict
-            kwargs passed to json.dump()
-
-        Returns
-        -------
-        str
-            Location of json file
-        """
-        idata_dict = _make_json_serializable(
-            self.to_dict(groups=groups, filter_groups=filter_groups)
-        )
-
-        with open(filename, "w", encoding="utf8") as file:
-            json.dump(idata_dict, file, **kwargs)
-
-        return filename
-
-    def to_dataframe(
-        self,
-        groups=None,
-        filter_groups=None,
-        var_names=None,
-        filter_vars=None,
-        include_coords=True,
-        include_index=True,
-        index_origin=None,
-    ):
-        """Convert InferenceData to a :class:`pandas.DataFrame` following xarray naming conventions.
-
-        This returns dataframe in a "wide" -format, where each item in ndimensional array is
-        unpacked. To access "tidy" -format, use xarray functionality found for each dataset.
-
-        In case of a multiple groups, function adds a group identification to the var name.
-
-        Data groups ("observed_data", "constant_data", "predictions_constant_data") are
-        skipped implicitly.
-
-        Raises TypeError if no valid groups are found.
-        Raises ValueError if no data are selected.
-
-        Parameters
-        ----------
-        groups: str or list of str, optional
-            Groups where the transformation is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups: {None, "like", "regex"}, optional, default=None
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        var_names : str or list of str, optional
-            Variables to be extracted. Prefix the variables by `~` when you want to exclude them.
-        filter_vars: {None, "like", "regex"}, optional
-            If `None` (default), interpret var_names as the real variables names. If "like",
-            interpret var_names as substrings of the real variables names. If "regex",
-            interpret var_names as regular expressions on the real variables names. A la
-            `pandas.filter`.
-            Like with plotting, sometimes it's easier to subset saying what to exclude
-            instead of what to include
-        include_coords: bool
-            Add coordinate values to column name (tuple).
-        include_index: bool
-            Add index information for multidimensional arrays.
-        index_origin: {0, 1}, optional
-            Starting index  for multidimensional objects. 0- or 1-based.
-            Defaults to rcParams["data.index_origin"].
-
-        Returns
-        -------
-        pandas.DataFrame
-            A pandas DataFrame containing all selected groups of InferenceData object.
-        """
-        # pylint: disable=too-many-nested-blocks
-        if not include_coords and not include_index:
-            raise TypeError("Both include_coords and include_index can not be False.")
-        if index_origin is None:
-            index_origin = rcParams["data.index_origin"]
-        if index_origin not in [0, 1]:
-            raise TypeError(f"index_origin must be 0 or 1, saw {index_origin}")
-
-        group_names = list(
-            filter(lambda x: "data" not in x, self._group_names(groups, filter_groups))
-        )
-
-        if not group_names:
-            raise TypeError(f"No valid groups found: {groups}")
-
-        dfs = {}
-        for group in group_names:
-            dataset = self[group]
-            group_var_names = _var_names(var_names, dataset, filter_vars, "ignore")
-            if (group_var_names is not None) and not group_var_names:
-                continue
-            if group_var_names is not None:
-                dataset = dataset[[var_name for var_name in group_var_names if var_name in dataset]]
-            df = None
-            coords_to_idx = {
-                name: dict(map(reversed, enumerate(dataset.coords[name].values, index_origin)))
-                for name in list(filter(lambda x: x not in ("chain", "draw"), dataset.coords))
-            }
-            for data_array in dataset.values():
-                dataframe = data_array.to_dataframe()
-                if list(filter(lambda x: x not in ("chain", "draw"), data_array.dims)):
-                    levels = [
-                        idx
-                        for idx, dim in enumerate(data_array.dims)
-                        if dim not in ("chain", "draw")
-                    ]
-                    dataframe = dataframe.unstack(level=levels)
-                    tuple_columns = []
-                    for name, *coords in dataframe.columns:
-                        if include_index:
-                            idxs = []
-                            for coordname, coorditem in zip(dataframe.columns.names[1:], coords):
-                                idxs.append(coords_to_idx[coordname][coorditem])
-                            if include_coords:
-                                tuple_columns.append(
-                                    (f"{name}[{','.join(map(str, idxs))}]", *coords)
-                                )
-                            else:
-                                tuple_columns.append(f"{name}[{','.join(map(str, idxs))}]")
-                        else:
-                            tuple_columns.append((name, *coords))
-
-                    dataframe.columns = tuple_columns
-                    dataframe.sort_index(axis=1, inplace=True)
-                if df is None:
-                    df = dataframe
-                    continue
-                df = df.join(dataframe, how="outer")
-            if df is not None:
-                df = df.reset_index()
-                dfs[group] = df
-        if not dfs:
-            raise ValueError("No data selected for the dataframe.")
-        if len(dfs) > 1:
-            for group, df in dfs.items():
-                df.columns = [
-                    (
-                        col
-                        if col in ("draw", "chain")
-                        else (group, *col) if isinstance(col, tuple) else (group, col)
-                    )
-                    for col in df.columns
-                ]
-            dfs, *dfs_tail = list(dfs.values())
-            for df in dfs_tail:
-                dfs = dfs.merge(df, how="outer", copy=False)
-        else:
-            (dfs,) = dfs.values()  # pylint: disable=unbalanced-dict-unpacking
-        return dfs
-
-    def to_zarr(self, store=None):
-        """Convert InferenceData to a :class:`zarr.hierarchy.Group`.
-
-        The zarr storage is using the same group names as the InferenceData.
-
-        Raises
-        ------
-        TypeError
-            If no valid store is found.
-
-        Parameters
-        ----------
-        store: zarr.storage i.e MutableMapping or str, optional
-            Zarr storage class or path to desired DirectoryStore.
-
-        Returns
-        -------
-        zarr.hierarchy.group
-            A zarr hierarchy group containing the InferenceData.
-
-        References
-        ----------
-        https://zarr.readthedocs.io/
-        """
-        try:
-            import zarr
-        except ImportError as err:
-            raise ImportError("'to_zarr' method needs Zarr (>=2.5.0,<3) installed.") from err
-        if version.parse(zarr.__version__) < version.parse("2.5.0"):
-            raise ImportError(
-                "Found zarr<2.5.0, please upgrade to a zarr (>=2.5.0,<3) to use 'to_zarr'"
-            )
-        if version.parse(zarr.__version__) >= version.parse("3.0.0.dev0"):
-            raise ImportError(
-                "Found zarr>=3, which is not supported by ArviZ. Instead, you can use "
-                "'dt = InferenceData.to_datatree' followed by 'dt.to_zarr()' "
-                "(needs xarray>=2024.11.0)"
-            )
-
-        # Check store type and create store if necessary
-        if store is None:
-            store = zarr.storage.TempStore(suffix="arviz")
-        elif isinstance(store, str):
-            store = zarr.storage.DirectoryStore(path=store)
-        elif not isinstance(store, MutableMapping):
-            raise TypeError(f"No valid store found: {store}")
-
-        groups = self.groups()
-
-        if not groups:
-            raise TypeError("No valid groups found!")
-
-        # order matters here, saving attrs after the groups will erase the groups.
-        if self.attrs:
-            xr.Dataset(attrs=self.attrs).to_zarr(store=store, mode="w")
-
-        for group in groups:
-            # Create zarr group in store with same group name
-            getattr(self, group).to_zarr(store=store, group=group, mode="w")
-
-        return zarr.open(store)  # Open store to get overarching group
-
-    @staticmethod
-    def from_zarr(store) -> "InferenceData":
-        """Initialize object from a zarr store or path.
-
-        Expects that the zarr store will have groups, each of which can be loaded by xarray.
-        By default, the datasets of the InferenceData object will be lazily loaded instead
-        of being loaded into memory. This
-        behaviour is regulated by the value of ``az.rcParams["data.load"]``.
-
-        Parameters
-        ----------
-        store: MutableMapping or zarr.hierarchy.Group or str.
-            Zarr storage class or path to desired Store.
-
-        Returns
-        -------
-        InferenceData object
-
-        References
-        ----------
-        https://zarr.readthedocs.io/
-        """
-        try:
-            import zarr
-        except ImportError as err:
-            raise ImportError("'from_zarr' method needs Zarr (>=2.5.0,<3) installed.") from err
-        if version.parse(zarr.__version__) < version.parse("2.5.0"):
-            raise ImportError(
-                "Found zarr<2.5.0, please upgrade to a zarr (>=2.5.0,<3) to use 'from_zarr'"
-            )
-        if version.parse(zarr.__version__) >= version.parse("3.0.0.dev0"):
-            raise ImportError(
-                "Found zarr>=3, which is not supported by ArviZ. Instead, you can use "
-                "'xarray.open_datatree' followed by 'arviz.InferenceData.from_datatree' "
-                "(needs xarray>=2024.11.0)"
-            )
-
-        # Check store type and create store if necessary
-        if isinstance(store, str):
-            store = zarr.storage.DirectoryStore(path=store)
-        elif isinstance(store, zarr.hierarchy.Group):
-            store = store.store
-        elif not isinstance(store, MutableMapping):
-            raise TypeError(f"No valid store found: {store}")
-
-        groups = {}
-        zarr_handle = zarr.open(store, mode="r")
-
-        # Open each group via xarray method
-        for key_group, _ in zarr_handle.groups():
-            with xr.open_zarr(store=store, group=key_group) as data:
-                groups[key_group] = data.load() if rcParams["data.load"] == "eager" else data
-
-        with xr.open_zarr(store=store) as root:
-            attrs = root.attrs
-
-        return InferenceData(attrs=attrs, **groups)
-
-    def __add__(self, other: "InferenceData") -> "InferenceData":
-        """Concatenate two InferenceData objects."""
-        return concat(self, other, copy=True, inplace=False)
-
-    def sel(
-        self: InferenceDataT,
-        groups: Optional[Union[str, List[str]]] = None,
-        filter_groups: Optional["Literal['like', 'regex']"] = None,
-        inplace: bool = False,
-        chain_prior: Optional[bool] = None,
-        **kwargs: Any,
-    ) -> Optional[InferenceDataT]:
-        """Perform an xarray selection on all groups.
-
-        Loops groups to perform Dataset.sel(key=item)
-        for every kwarg if key is a dimension of the dataset.
-        One example could be performing a burn in cut on the InferenceData object
-        or discarding a chain. The selection is performed on all relevant groups (like
-        posterior, prior, sample stats) while non relevant groups like observed data are
-        omitted. See :meth:`xarray.Dataset.sel <xarray:xarray.Dataset.sel>`
-
-        Parameters
-        ----------
-        groups : str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional, default=None
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-        chain_prior : bool, optional, deprecated
-            If ``False``, do not select prior related groups using ``chain`` dim.
-            Otherwise, use selection on ``chain`` if present. Default=False
-        kwargs : dict, optional
-            It must be accepted by Dataset.sel().
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object by default.
-            When `inplace==True` perform selection in-place and return `None`
-
-        Examples
-        --------
-        Use ``sel`` to discard one chain of the InferenceData object. We first check the
-        dimensions of the original object:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("centered_eight")
-            idata
-
-        In order to remove the third chain:
-
-        .. jupyter-execute::
-
-            idata_subset = idata.sel(chain=[0, 1, 3], groups="posterior_groups")
-            idata_subset
-
-        See Also
-        --------
-        xarray.Dataset.sel :
-            Returns a new dataset with each array indexed by tick labels along the specified
-            dimension(s).
-        isel : Returns a new dataset with each array indexed along the specified dimension(s).
-        """
-        if chain_prior is not None:
-            warnings.warn(
-                "chain_prior has been deprecated. Use groups argument and "
-                "rcParams['data.metagroups'] instead.",
-                DeprecationWarning,
-            )
-        else:
-            chain_prior = False
-        group_names = self._group_names(groups, filter_groups)
-
-        out = self if inplace else deepcopy(self)
-        for group in group_names:
-            dataset = getattr(self, group)
-            valid_keys = set(kwargs.keys()).intersection(dataset.dims)
-            if not chain_prior and "prior" in group:
-                valid_keys -= {"chain"}
-            dataset = dataset.sel(**{key: kwargs[key] for key in valid_keys})
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def isel(
-        self: InferenceDataT,
-        groups: Optional[Union[str, List[str]]] = None,
-        filter_groups: Optional["Literal['like', 'regex']"] = None,
-        inplace: bool = False,
-        **kwargs: Any,
-    ) -> Optional[InferenceDataT]:
-        """Perform an xarray selection on all groups.
-
-        Loops groups to perform Dataset.isel(key=item)
-        for every kwarg if key is a dimension of the dataset.
-        One example could be performing a burn in cut on the InferenceData object
-        or discarding a chain. The selection is performed on all relevant groups (like
-        posterior, prior, sample stats) while non relevant groups like observed data are
-        omitted. See :meth:`xarray:xarray.Dataset.isel`
-
-        Parameters
-        ----------
-        groups : str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-        kwargs : dict, optional
-            It must be accepted by :meth:`xarray:xarray.Dataset.isel`.
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object by default.
-            When `inplace==True` perform selection in-place and return `None`
-
-        Examples
-        --------
-        Use ``isel`` to discard one chain of the InferenceData object. We first check the
-        dimensions of the original object:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("centered_eight")
-            idata
-
-        In order to remove the third chain:
-
-        .. jupyter-execute::
-
-            idata_subset = idata.isel(chain=[0, 1, 3], groups="posterior_groups")
-            idata_subset
-
-        You can expand the groups and coords in each group to see how now only the chains 0, 1 and
-        3 are present.
-
-        See Also
-        --------
-        xarray.Dataset.isel :
-            Returns a new dataset with each array indexed along the specified dimension(s).
-        sel :
-            Returns a new dataset with each array indexed by tick labels along the specified
-            dimension(s).
-        """
-        group_names = self._group_names(groups, filter_groups)
-
-        out = self if inplace else deepcopy(self)
-        for group in group_names:
-            dataset = getattr(self, group)
-            valid_keys = set(kwargs.keys()).intersection(dataset.dims)
-            dataset = dataset.isel(**{key: kwargs[key] for key in valid_keys})
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def stack(
-        self,
-        dimensions=None,
-        groups=None,
-        filter_groups=None,
-        inplace=False,
-        **kwargs,
-    ):
-        """Perform an xarray stacking on all groups.
-
-        Stack any number of existing dimensions into a single new dimension.
-        Loops groups to perform Dataset.stack(key=value)
-        for every kwarg if value is a dimension of the dataset.
-        The selection is performed on all relevant groups (like
-        posterior, prior, sample stats) while non relevant groups like observed data are
-        omitted. See :meth:`xarray:xarray.Dataset.stack`
-
-        Parameters
-        ----------
-        dimensions : dict, optional
-            Names of new dimensions, and the existing dimensions that they replace.
-        groups: str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-        kwargs : dict, optional
-            It must be accepted by :meth:`xarray:xarray.Dataset.stack`.
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object by default.
-            When `inplace==True` perform selection in-place and return `None`
-
-        Examples
-        --------
-        Use ``stack`` to stack any number of existing dimensions into a single new dimension.
-        We first check the original object:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("rugby")
-            idata
-
-        In order to stack two dimensions ``chain`` and ``draw`` to ``sample``, we can use:
-
-        .. jupyter-execute::
-
-            idata.stack(sample=["chain", "draw"], inplace=True)
-            idata
-
-        We can also take the example of custom InferenceData object and perform stacking. We first
-        check the original object:
-
-        .. jupyter-execute::
-
-            import numpy as np
-            datadict = {
-                "a": np.random.randn(100),
-                "b": np.random.randn(1, 100, 10),
-                "c": np.random.randn(1, 100, 3, 4),
-            }
-            coords = {
-                "c1": np.arange(3),
-                "c99": np.arange(4),
-                "b1": np.arange(10),
-            }
-            dims = {"c": ["c1", "c99"], "b": ["b1"]}
-            idata = az.from_dict(
-                posterior=datadict, posterior_predictive=datadict, coords=coords, dims=dims
-            )
-            idata
-
-        In order to stack two dimensions ``c1`` and ``c99`` to ``z``, we can use:
-
-        .. jupyter-execute::
-
-            idata.stack(z=["c1", "c99"], inplace=True)
-            idata
-
-        See Also
-        --------
-        xarray.Dataset.stack : Stack any number of existing dimensions into a single new dimension.
-        unstack : Perform an xarray unstacking on all groups of InferenceData object.
-        """
-        groups = self._group_names(groups, filter_groups)
-
-        dimensions = {} if dimensions is None else dimensions
-        dimensions.update(kwargs)
-        out = self if inplace else deepcopy(self)
-        for group in groups:
-            dataset = getattr(self, group)
-            kwarg_dict = {}
-            for key, value in dimensions.items():
-                try:
-                    if not set(value).difference(dataset.dims):
-                        kwarg_dict[key] = value
-                except TypeError:
-                    kwarg_dict[key] = value
-            dataset = dataset.stack(**kwarg_dict)
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def unstack(self, dim=None, groups=None, filter_groups=None, inplace=False):
-        """Perform an xarray unstacking on all groups.
-
-        Unstack existing dimensions corresponding to MultiIndexes into multiple new dimensions.
-        Loops groups to perform Dataset.unstack(key=value).
-        The selection is performed on all relevant groups (like posterior, prior,
-        sample stats) while non relevant groups like observed data are omitted.
-        See :meth:`xarray:xarray.Dataset.unstack`
-
-        Parameters
-        ----------
-        dim : Hashable or iterable of Hashable, optional
-            Dimension(s) over which to unstack. By default unstacks all MultiIndexes.
-        groups : str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object by default.
-            When `inplace==True` perform selection in place and return `None`
-
-        Examples
-        --------
-        Use ``unstack`` to unstack existing dimensions corresponding to MultiIndexes into
-        multiple new dimensions. We first stack two dimensions ``c1`` and ``c99`` to ``z``:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            import numpy as np
-            datadict = {
-                "a": np.random.randn(100),
-                "b": np.random.randn(1, 100, 10),
-                "c": np.random.randn(1, 100, 3, 4),
-            }
-            coords = {
-                "c1": np.arange(3),
-                "c99": np.arange(4),
-                "b1": np.arange(10),
-            }
-            dims = {"c": ["c1", "c99"], "b": ["b1"]}
-            idata = az.from_dict(
-                posterior=datadict, posterior_predictive=datadict, coords=coords, dims=dims
-            )
-            idata.stack(z=["c1", "c99"], inplace=True)
-            idata
-
-        In order to unstack the dimension ``z``, we use:
-
-        .. jupyter-execute::
-
-            idata.unstack(inplace=True)
-            idata
-
-        See Also
-        --------
-        xarray.Dataset.unstack :
-            Unstack existing dimensions corresponding to MultiIndexes into multiple new dimensions.
-        stack : Perform an xarray stacking on all groups of InferenceData object.
-        """
-        groups = self._group_names(groups, filter_groups)
-        if isinstance(dim, str):
-            dim = [dim]
-
-        out = self if inplace else deepcopy(self)
-        for group in groups:
-            dataset = getattr(self, group)
-            valid_dims = set(dim).intersection(dataset.dims) if dim is not None else dim
-            dataset = dataset.unstack(dim=valid_dims)
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def rename(self, name_dict=None, groups=None, filter_groups=None, inplace=False):
-        """Perform xarray renaming of variable and dimensions on all groups.
-
-        Loops groups to perform Dataset.rename(name_dict)
-        for every key in name_dict if key is a dimension/data_vars of the dataset.
-        The renaming is performed on all relevant groups (like
-        posterior, prior, sample stats) while non relevant groups like observed data are
-        omitted. See :meth:`xarray:xarray.Dataset.rename`
-
-        Parameters
-        ----------
-        name_dict : dict
-            Dictionary whose keys are current variable or dimension names
-            and whose values are the desired names.
-        groups : str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object by default.
-            When `inplace==True` perform renaming in-place and return `None`
-
-        Examples
-        --------
-        Use ``rename`` to renaming of variable and dimensions on all groups of the InferenceData
-        object. We first check the original object:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("rugby")
-            idata
-
-        In order to rename the dimensions and variable, we use:
-
-        .. jupyter-execute::
-
-            idata.rename({"team": "team_new", "match":"match_new"}, inplace=True)
-            idata
-
-        See Also
-        --------
-        xarray.Dataset.rename : Returns a new object with renamed variables and dimensions.
-        rename_vars :
-            Perform xarray renaming of variable or coordinate names on all groups of an
-            InferenceData object.
-        rename_dims : Perform xarray renaming of dimensions on all groups of InferenceData object.
-        """
-        groups = self._group_names(groups, filter_groups)
-        if "chain" in name_dict.keys() or "draw" in name_dict.keys():
-            raise KeyError("'chain' or 'draw' dimensions can't be renamed")
-        out = self if inplace else deepcopy(self)
-
-        for group in groups:
-            dataset = getattr(self, group)
-            expected_keys = list(dataset.data_vars) + list(dataset.dims)
-            valid_keys = set(name_dict.keys()).intersection(expected_keys)
-            dataset = dataset.rename({key: name_dict[key] for key in valid_keys})
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def rename_vars(self, name_dict=None, groups=None, filter_groups=None, inplace=False):
-        """Perform xarray renaming of variable or coordinate names on all groups.
-
-        Loops groups to perform Dataset.rename_vars(name_dict)
-        for every key in name_dict if key is a variable or coordinate names of the dataset.
-        The renaming is performed on all relevant groups (like
-        posterior, prior, sample stats) while non relevant groups like observed data are
-        omitted. See :meth:`xarray:xarray.Dataset.rename_vars`
-
-        Parameters
-        ----------
-        name_dict : dict
-            Dictionary whose keys are current variable or coordinate names
-            and whose values are the desired names.
-        groups : str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object with renamed variables including coordinates by default.
-            When `inplace==True` perform renaming in-place and return `None`
-
-        Examples
-        --------
-        Use ``rename_vars`` to renaming of variable and coordinates on all groups of the
-        InferenceData object. We first check the data variables of original object:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("rugby")
-            idata
-
-        In order to rename the data variables, we use:
-
-        .. jupyter-execute::
-
-            idata.rename_vars({"home": "home_new"}, inplace=True)
-            idata
-
-        See Also
-        --------
-        xarray.Dataset.rename_vars :
-            Returns a new object with renamed variables including coordinates.
-        rename :
-            Perform xarray renaming of variable and dimensions on all groups of an InferenceData
-            object.
-        rename_dims : Perform xarray renaming of dimensions on all groups of InferenceData object.
-        """
-        groups = self._group_names(groups, filter_groups)
-
-        out = self if inplace else deepcopy(self)
-        for group in groups:
-            dataset = getattr(self, group)
-            valid_keys = set(name_dict.keys()).intersection(dataset.data_vars)
-            dataset = dataset.rename_vars({key: name_dict[key] for key in valid_keys})
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def rename_dims(self, name_dict=None, groups=None, filter_groups=None, inplace=False):
-        """Perform xarray renaming of dimensions on all groups.
-
-        Loops groups to perform Dataset.rename_dims(name_dict)
-        for every key in name_dict if key is a dimension of the dataset.
-        The renaming is performed on all relevant groups (like
-        posterior, prior, sample stats) while non relevant groups like observed data are
-        omitted. See :meth:`xarray:xarray.Dataset.rename_dims`
-
-        Parameters
-        ----------
-        name_dict : dict
-            Dictionary whose keys are current dimension names and whose values are the desired
-            names.
-        groups : str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object with renamed dimension by default.
-            When `inplace==True` perform renaming in-place and return `None`
-
-        Examples
-        --------
-        Use ``rename_dims`` to renaming of dimensions on all groups of the InferenceData
-        object. We first check the dimensions of original object:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("rugby")
-            idata
-
-        In order to rename the dimensions, we use:
-
-        .. jupyter-execute::
-
-            idata.rename_dims({"team": "team_new"}, inplace=True)
-            idata
-
-        See Also
-        --------
-        xarray.Dataset.rename_dims : Returns a new object with renamed dimensions only.
-        rename :
-            Perform xarray renaming of variable and dimensions on all groups of an InferenceData
-            object.
-        rename_vars :
-            Perform xarray renaming of variable or coordinate names on all groups of an
-            InferenceData object.
-        """
-        groups = self._group_names(groups, filter_groups)
-        if "chain" in name_dict.keys() or "draw" in name_dict.keys():
-            raise KeyError("'chain' or 'draw' dimensions can't be renamed")
-
-        out = self if inplace else deepcopy(self)
-        for group in groups:
-            dataset = getattr(self, group)
-            valid_keys = set(name_dict.keys()).intersection(dataset.dims)
-            dataset = dataset.rename_dims({key: name_dict[key] for key in valid_keys})
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def add_groups(
-        self, group_dict=None, coords=None, dims=None, warn_on_custom_groups=False, **kwargs
-    ):
-        """Add new groups to InferenceData object.
-
-        Parameters
-        ----------
-        group_dict : dict of {str : dict or xarray.Dataset}, optional
-            Groups to be added
-        coords : dict of {str : array_like}, optional
-            Coordinates for the dataset
-        dims : dict of {str : list of str}, optional
-            Dimensions of each variable. The keys are variable names, values are lists of
-            coordinates.
-        warn_on_custom_groups : bool, default False
-            Emit a warning when custom groups are present in the InferenceData.
-            "custom group" means any group whose name isn't defined in :ref:`schema`
-        kwargs : dict, optional
-            The keyword arguments form of group_dict. One of group_dict or kwargs must be provided.
-
-        Examples
-        --------
-        Add a ``log_likelihood`` group to the "rugby" example InferenceData after loading.
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("rugby")
-            del idata.log_likelihood
-            idata2 = idata.copy()
-            post = idata.posterior
-            obs = idata.observed_data
-            idata
-
-        Knowing the model, we can compute it manually. In this case however,
-        we will generate random samples with the right shape.
-
-        .. jupyter-execute::
-
-            import numpy as np
-            rng = np.random.default_rng(73)
-            ary = rng.normal(size=(post.sizes["chain"], post.sizes["draw"], obs.sizes["match"]))
-            idata.add_groups(
-                log_likelihood={"home_points": ary},
-                dims={"home_points": ["match"]},
-            )
-            idata
-
-        This is fine if we have raw data, but a bit inconvenient if we start with labeled
-        data already. Why provide dims and coords manually again?
-        Let's generate a fake log likelihood (doesn't match the model but it serves just
-        the same for illustration purposes here) working from the posterior and
-        observed_data groups manually:
-
-        .. jupyter-execute::
-
-            import xarray as xr
-            from xarray_einstats.stats import XrDiscreteRV
-            from scipy.stats import poisson
-            dist = XrDiscreteRV(poisson, np.exp(post["atts"]))
-            log_lik = dist.logpmf(obs["home_points"]).to_dataset(name="home_points")
-            idata2.add_groups({"log_likelihood": log_lik})
-            idata2
-
-        Note that in the first example we have used the ``kwargs`` argument
-        and in the second we have used the ``group_dict`` one.
-
-        See Also
-        --------
-        extend : Extend InferenceData with groups from another InferenceData.
-        concat : Concatenate InferenceData objects.
-        """
-        group_dict = either_dict_or_kwargs(group_dict, kwargs, "add_groups")
-        if not group_dict:
-            raise ValueError("One of group_dict or kwargs must be provided.")
-        repeated_groups = [group for group in group_dict.keys() if group in self._groups]
-        if repeated_groups:
-            raise ValueError(f"{repeated_groups} group(s) already exists.")
-        for group, dataset in group_dict.items():
-            if warn_on_custom_groups and group not in SUPPORTED_GROUPS_ALL:
-                warnings.warn(
-                    f"The group {group} is not defined in the InferenceData scheme",
-                    UserWarning,
-                )
-            if dataset is None:
-                continue
-            elif isinstance(dataset, dict):
-                if (
-                    group in ("observed_data", "constant_data", "predictions_constant_data")
-                    or group not in SUPPORTED_GROUPS_ALL
-                ):
-                    warnings.warn(
-                        "the default dims 'chain' and 'draw' will be added automatically",
-                        UserWarning,
-                    )
-                dataset = dict_to_dataset(dataset, coords=coords, dims=dims)
-            elif isinstance(dataset, xr.DataArray):
-                if dataset.name is None:
-                    dataset.name = "x"
-                dataset = dataset.to_dataset()
-            elif not isinstance(dataset, xr.Dataset):
-                raise ValueError(
-                    "Arguments to add_groups() must be xr.Dataset, xr.Dataarray or dicts\
-                    (argument '{}' was type '{}')".format(
-                        group, type(dataset)
-                    )
-                )
-            if dataset:
-                setattr(self, group, dataset)
-                if group.startswith(WARMUP_TAG):
-                    supported_order = [
-                        key for key in SUPPORTED_GROUPS_ALL if key in self._groups_warmup
-                    ]
-                    if (supported_order == self._groups_warmup) and (group in SUPPORTED_GROUPS_ALL):
-                        group_order = [
-                            key
-                            for key in SUPPORTED_GROUPS_ALL
-                            if key in self._groups_warmup + [group]
-                        ]
-                        group_idx = group_order.index(group)
-                        self._groups_warmup.insert(group_idx, group)
-                    else:
-                        self._groups_warmup.append(group)
-                else:
-                    supported_order = [key for key in SUPPORTED_GROUPS_ALL if key in self._groups]
-                    if (supported_order == self._groups) and (group in SUPPORTED_GROUPS_ALL):
-                        group_order = [
-                            key for key in SUPPORTED_GROUPS_ALL if key in self._groups + [group]
-                        ]
-                        group_idx = group_order.index(group)
-                        self._groups.insert(group_idx, group)
-                    else:
-                        self._groups.append(group)
-
-    def extend(self, other, join="left", warn_on_custom_groups=False):
-        """Extend InferenceData with groups from another InferenceData.
-
-        Parameters
-        ----------
-        other : InferenceData
-            InferenceData to be added
-        join : {'left', 'right'}, default 'left'
-            Defines how the two decide which group to keep when the same group is
-            present in both objects. 'left' will discard the group in ``other`` whereas 'right'
-            will keep the group in ``other`` and discard the one in ``self``.
-        warn_on_custom_groups : bool, default False
-            Emit a warning when custom groups are present in the InferenceData.
-            "custom group" means any group whose name isn't defined in :ref:`schema`
-
-        Examples
-        --------
-        Take two InferenceData objects, and extend the first with the groups it doesn't have
-        but are present in the 2nd InferenceData object.
-
-        First InferenceData:
-
-        .. jupyter-execute::
-
-            import arviz as az
-            idata = az.load_arviz_data("radon")
-
-        Second InferenceData:
-
-        .. jupyter-execute::
-
-            other_idata = az.load_arviz_data("rugby")
-
-        Call the ``extend`` method:
-
-        .. jupyter-execute::
-
-            idata.extend(other_idata)
-            idata
-
-        See how now the first InferenceData has more groups, with the data from the
-        second one, but the groups it originally had have not been modified,
-        even if also present in the second InferenceData.
-
-        See Also
-        --------
-        add_groups : Add new groups to InferenceData object.
-        concat : Concatenate InferenceData objects.
-
-        """
-        if not isinstance(other, InferenceData):
-            raise ValueError("Extending is possible between two InferenceData objects only.")
-        if join not in ("left", "right"):
-            raise ValueError(f"join must be either 'left' or 'right', found {join}")
-        for group in other._groups_all:  # pylint: disable=protected-access
-            if hasattr(self, group) and join == "left":
-                continue
-            if warn_on_custom_groups and group not in SUPPORTED_GROUPS_ALL:
-                warnings.warn(
-                    f"{group} group is not defined in the InferenceData scheme", UserWarning
-                )
-            dataset = getattr(other, group)
-            setattr(self, group, dataset)
-            if group.startswith(WARMUP_TAG):
-                if group not in self._groups_warmup:
-                    supported_order = [
-                        key for key in SUPPORTED_GROUPS_ALL if key in self._groups_warmup
-                    ]
-                    if (supported_order == self._groups_warmup) and (group in SUPPORTED_GROUPS_ALL):
-                        group_order = [
-                            key
-                            for key in SUPPORTED_GROUPS_ALL
-                            if key in self._groups_warmup + [group]
-                        ]
-                        group_idx = group_order.index(group)
-                        self._groups_warmup.insert(group_idx, group)
-                    else:
-                        self._groups_warmup.append(group)
-            elif group not in self._groups:
-                supported_order = [key for key in SUPPORTED_GROUPS_ALL if key in self._groups]
-                if (supported_order == self._groups) and (group in SUPPORTED_GROUPS_ALL):
-                    group_order = [
-                        key for key in SUPPORTED_GROUPS_ALL if key in self._groups + [group]
-                    ]
-                    group_idx = group_order.index(group)
-                    self._groups.insert(group_idx, group)
-                else:
-                    self._groups.append(group)
-
-    set_index = _extend_xr_method(xr.Dataset.set_index, see_also="reset_index")
-    get_index = _extend_xr_method(xr.Dataset.get_index)
-    reset_index = _extend_xr_method(xr.Dataset.reset_index, see_also="set_index")
-    set_coords = _extend_xr_method(xr.Dataset.set_coords, see_also="reset_coords")
-    reset_coords = _extend_xr_method(xr.Dataset.reset_coords, see_also="set_coords")
-    assign = _extend_xr_method(xr.Dataset.assign)
-    assign_coords = _extend_xr_method(xr.Dataset.assign_coords)
-    sortby = _extend_xr_method(xr.Dataset.sortby)
-    chunk = _extend_xr_method(xr.Dataset.chunk)
-    unify_chunks = _extend_xr_method(xr.Dataset.unify_chunks)
-    load = _extend_xr_method(xr.Dataset.load)
-    compute = _extend_xr_method(xr.Dataset.compute)
-    persist = _extend_xr_method(xr.Dataset.persist)
-    quantile = _extend_xr_method(xr.Dataset.quantile)
-    close = _extend_xr_method(xr.Dataset.close)
-
-    # The following lines use methods on xr.Dataset that are dynamically defined and attached.
-    # As a result mypy cannot see them, so we have to suppress the resulting mypy errors.
-    mean = _extend_xr_method(xr.Dataset.mean, see_also="median")  # type: ignore[attr-defined]
-    median = _extend_xr_method(xr.Dataset.median, see_also="mean")  # type: ignore[attr-defined]
-    min = _extend_xr_method(xr.Dataset.min, see_also=["max", "sum"])  # type: ignore[attr-defined]
-    max = _extend_xr_method(xr.Dataset.max, see_also=["min", "sum"])  # type: ignore[attr-defined]
-    cumsum = _extend_xr_method(xr.Dataset.cumsum, see_also="sum")  # type: ignore[attr-defined]
-    sum = _extend_xr_method(xr.Dataset.sum, see_also="cumsum")  # type: ignore[attr-defined]
-
-    def _group_names(
-        self,
-        groups: Optional[Union[str, List[str]]],
-        filter_groups: Optional["Literal['like', 'regex']"] = None,
-    ) -> List[str]:
-        """Handle expansion of group names input across arviz.
-
-        Parameters
-        ----------
-        groups: str, list of str or None
-            group or metagroup names.
-        idata: xarray.Dataset
-            Posterior data in an xarray
-        filter_groups: {None, "like", "regex"}, optional, default=None
-            If `None` (default), interpret groups as the real group or metagroup names.
-            If "like", interpret groups as substrings of the real group or metagroup names.
-            If "regex", interpret groups as regular expressions on the real group or
-            metagroup names. A la `pandas.filter`.
-
-        Returns
-        -------
-        groups: list
-        """
-        if filter_groups not in {None, "like", "regex"}:
-            raise ValueError(
-                f"'filter_groups' can only be None, 'like', or 'regex', got: '{filter_groups}'"
-            )
-
-        all_groups = self._groups_all
-        if groups is None:
-            return all_groups
-        if isinstance(groups, str):
-            groups = [groups]
-        sel_groups = []
-        metagroups = rcParams["data.metagroups"]
-        for group in groups:
-            if group[0] == "~":
-                sel_groups.extend(
-                    [f"~{item}" for item in metagroups[group[1:]] if item in all_groups]
-                    if group[1:] in metagroups
-                    else [group]
-                )
-            else:
-                sel_groups.extend(
-                    [item for item in metagroups[group] if item in all_groups]
-                    if group in metagroups
-                    else [group]
-                )
-
-        try:
-            group_names = _subset_list(sel_groups, all_groups, filter_items=filter_groups)
-        except KeyError as err:
-            msg = " ".join(("groups:", f"{err}", "in InferenceData"))
-            raise KeyError(msg) from err
-        return group_names
-
-    def map(self, fun, groups=None, filter_groups=None, inplace=False, args=None, **kwargs):
-        """Apply a function to multiple groups.
-
-        Applies ``fun`` groupwise to the selected ``InferenceData`` groups and overwrites the
-        group with the result of the function.
-
-        Parameters
-        ----------
-        fun : callable
-            Function to be applied to each group. Assumes the function is called as
-            ``fun(dataset, *args, **kwargs)``.
-        groups : str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        filter_groups : {None, "like", "regex"}, optional
-            If `None` (default), interpret var_names as the real variables names. If "like",
-            interpret var_names as substrings of the real variables names. If "regex",
-            interpret var_names as regular expressions on the real variables names. A la
-            `pandas.filter`.
-        inplace : bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-        args : array_like, optional
-            Positional arguments passed to ``fun``.
-        **kwargs : mapping, optional
-            Keyword arguments passed to ``fun``.
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object by default.
-            When `inplace==True` perform selection in place and return `None`
-
-        Examples
-        --------
-        Shift observed_data, prior_predictive and posterior_predictive.
-
-        .. jupyter-execute::
-
-            import arviz as az
-            import numpy as np
-            idata = az.load_arviz_data("non_centered_eight")
-            idata_shifted_obs = idata.map(lambda x: x + 3, groups="observed_vars")
-            idata_shifted_obs
-
-        Rename and update the coordinate values in both posterior and prior groups.
-
-        .. jupyter-execute::
-
-            idata = az.load_arviz_data("radon")
-            idata = idata.map(
-                lambda ds: ds.rename({"g_coef": "uranium_coefs"}).assign(
-                    uranium_coefs=["intercept", "u_slope"]
-                ),
-                groups=["posterior", "prior"]
-            )
-            idata
-
-        Add extra coordinates to all groups containing observed variables
-
-        .. jupyter-execute::
-
-            idata = az.load_arviz_data("rugby")
-            home_team, away_team = np.array([
-                m.split() for m in idata.observed_data.match.values
-            ]).T
-            idata = idata.map(
-                lambda ds, **kwargs: ds.assign_coords(**kwargs),
-                groups="observed_vars",
-                home_team=("match", home_team),
-                away_team=("match", away_team),
-            )
-            idata
-
-        """
-        if args is None:
-            args = []
-        groups = self._group_names(groups, filter_groups)
-
-        out = self if inplace else deepcopy(self)
-        for group in groups:
-            dataset = getattr(self, group)
-            dataset = fun(dataset, *args, **kwargs)
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def _wrap_xarray_method(
-        self, method, groups=None, filter_groups=None, inplace=False, args=None, **kwargs
-    ):
-        """Extend and xarray.Dataset method to InferenceData object.
-
-        Parameters
-        ----------
-        method: str
-            Method to be extended. Must be a ``xarray.Dataset`` method.
-        groups: str or list of str, optional
-            Groups where the selection is to be applied. Can either be group names
-            or metagroup names.
-        inplace: bool, optional
-            If ``True``, modify the InferenceData object inplace,
-            otherwise, return the modified copy.
-        **kwargs: mapping, optional
-            Keyword arguments passed to the xarray Dataset method.
-
-        Returns
-        -------
-        InferenceData
-            A new InferenceData object by default.
-            When `inplace==True` perform selection in place and return `None`
-
-        Examples
-        --------
-        Compute the mean of `posterior_groups`:
-
-        .. ipython::
-
-            In [1]: import arviz as az
-               ...: idata = az.load_arviz_data("non_centered_eight")
-               ...: idata_means = idata._wrap_xarray_method("mean", groups="latent_vars")
-               ...: print(idata_means.posterior)
-               ...: print(idata_means.observed_data)
-
-        .. ipython::
-
-            In [1]: idata_stack = idata._wrap_xarray_method(
-               ...:     "stack",
-               ...:     groups=["posterior_groups", "prior_groups"],
-               ...:     sample=["chain", "draw"]
-               ...: )
-               ...: print(idata_stack.posterior)
-               ...: print(idata_stack.prior)
-               ...: print(idata_stack.observed_data)
-
-        """
-        if args is None:
-            args = []
-        groups = self._group_names(groups, filter_groups)
-
-        method = getattr(xr.Dataset, method)
-
-        out = self if inplace else deepcopy(self)
-        for group in groups:
-            dataset = getattr(self, group)
-            dataset = method(dataset, *args, **kwargs)
-            setattr(out, group, dataset)
-        if inplace:
-            return None
-        else:
-            return out
-
-    def copy(self) -> "InferenceData":
-        """Return a fresh copy of the ``InferenceData`` object."""
-        return deepcopy(self)
-
-
-@overload
-def concat(
-    *args,
-    dim: Optional[str] = None,
-    copy: bool = True,
-    inplace: "Literal[True]",
-    reset_dim: bool = True,
-) -> None: ...
-
-
-@overload
-def concat(
-    *args,
-    dim: Optional[str] = None,
-    copy: bool = True,
-    inplace: "Literal[False]",
-    reset_dim: bool = True,
-) -> InferenceData: ...
-
-
-@overload
-def concat(
-    ids: Iterable[InferenceData],
-    dim: Optional[str] = None,
-    *,
-    copy: bool = True,
-    inplace: "Literal[False]",
-    reset_dim: bool = True,
-) -> InferenceData: ...
-
-
-@overload
-def concat(
-    ids: Iterable[InferenceData],
-    dim: Optional[str] = None,
-    *,
-    copy: bool = True,
-    inplace: "Literal[True]",
-    reset_dim: bool = True,
-) -> None: ...
-
-
-@overload
-def concat(
-    ids: Iterable[InferenceData],
-    dim: Optional[str] = None,
-    *,
-    copy: bool = True,
-    inplace: bool = False,
-    reset_dim: bool = True,
-) -> Optional[InferenceData]: ...
-
-
-# pylint: disable=protected-access, inconsistent-return-statements
-def concat(*args, dim=None, copy=True, inplace=False, reset_dim=True):
-    """Concatenate InferenceData objects.
-
-    Concatenates over `group`, `chain` or `draw`.
-    By default concatenates over unique groups.
-    To concatenate over `chain` or `draw` function
-    needs identical groups and variables.
-
-    The `variables` in the `data` -group are merged if `dim` are not found.
-
-
-    Parameters
-    ----------
-    *args : InferenceData
-        Variable length InferenceData list or
-        Sequence of InferenceData.
-    dim : str, optional
-        If defined, concatenated over the defined dimension.
-        Dimension which is concatenated. If None, concatenates over
-        unique groups.
-    copy : bool
-        If True, groups are copied to the new InferenceData object.
-        Used only if `dim` is None.
-    inplace : bool
-        If True, merge args to first object.
-    reset_dim : bool
-        Valid only if dim is not None.
-
-    Returns
-    -------
-    InferenceData
-        A new InferenceData object by default.
-        When `inplace==True` merge args to first arg and return `None`
-
-    See Also
-    --------
-    add_groups : Add new groups to InferenceData object.
-    extend : Extend InferenceData with groups from another InferenceData.
-
-    Examples
-    --------
-    Use ``concat`` method to concatenate InferenceData objects. This will concatenates over
-    unique groups by default. We first create an ``InferenceData`` object:
-
-    .. ipython::
-
-        In [1]: import arviz as az
-           ...: import numpy as np
-           ...: data = {
-           ...:     "a": np.random.normal(size=(4, 100, 3)),
-           ...:     "b": np.random.normal(size=(4, 100)),
-           ...: }
-           ...: coords = {"a_dim": ["x", "y", "z"]}
-           ...: dataA = az.from_dict(data, coords=coords, dims={"a": ["a_dim"]})
-           ...: dataA
-
-    We have created an ``InferenceData`` object with default group 'posterior'. Now, we will
-    create another ``InferenceData`` object:
-
-    .. ipython::
-
-        In [1]: dataB = az.from_dict(prior=data, coords=coords, dims={"a": ["a_dim"]})
-           ...: dataB
-
-    We have created another ``InferenceData`` object with group 'prior'. Now, we will concatenate
-    these two ``InferenceData`` objects:
-
-    .. ipython::
-
-        In [1]: az.concat(dataA, dataB)
-
-    Now, we will concatenate over chain (or draw). It requires identical groups and variables.
-    Here we are concatenating two identical ``InferenceData`` objects over dimension chain:
-
-    .. ipython::
-
-        In [1]: az.concat(dataA, dataA, dim="chain")
-
-    It will create an ``InferenceData`` with the original group 'posterior'. In similar way,
-    we can also concatenate over draws.
-
-    """
-    # pylint: disable=undefined-loop-variable, too-many-nested-blocks
-    if len(args) == 0:
-        if inplace:
-            return
-        return InferenceData()
-
-    if len(args) == 1 and isinstance(args[0], Sequence):
-        args = args[0]
-
-    # assert that all args are InferenceData
-    for i, arg in enumerate(args):
-        if not isinstance(arg, InferenceData):
-            raise TypeError(
-                "Concatenating is supported only"
-                "between InferenceData objects. Input arg {} is {}".format(i, type(arg))
-            )
-
-    if dim is not None and dim.lower() not in {"group", "chain", "draw"}:
-        msg = f'Invalid `dim`: {dim}. Valid `dim` are {{"group", "chain", "draw"}}'
-        raise TypeError(msg)
-    dim = dim.lower() if dim is not None else dim
-
-    if len(args) == 1 and isinstance(args[0], InferenceData):
-        if inplace:
-            return None
-        else:
-            if copy:
-                return deepcopy(args[0])
-            else:
-                return args[0]
-
-    current_time = datetime.datetime.now(datetime.timezone.utc).isoformat()
-    combined_attr = defaultdict(list)
-    for idata in args:
-        for key, val in idata.attrs.items():
-            combined_attr[key].append(val)
-
-    for key, val in combined_attr.items():
-        all_same = True
-        for indx in range(len(val) - 1):
-            if val[indx] != val[indx + 1]:
-                all_same = False
-                break
-        if all_same:
-            combined_attr[key] = val[0]
-    if inplace:
-        setattr(args[0], "_attrs", dict(combined_attr))
-
-    if not inplace:
-        # Keep order for python 3.5
-        inference_data_dict = OrderedDict()
-
-    if dim is None:
-        arg0 = args[0]
-        arg0_groups = ccopy(arg0._groups_all)
-        args_groups = {}
-        # check if groups are independent
-        # Concat over unique groups
-        for arg in args[1:]:
-            for group in arg._groups_all:
-                if group in args_groups or group in arg0_groups:
-                    msg = (
-                        "Concatenating overlapping groups is not supported unless `dim` is defined."
-                        " Valid dimensions are `chain` and `draw`. Alternatively, use extend to"
-                        " combine InferenceData with overlapping groups"
-                    )
-                    raise TypeError(msg)
-                group_data = getattr(arg, group)
-                args_groups[group] = deepcopy(group_data) if copy else group_data
-        # add arg0 to args_groups if inplace is False
-        # otherwise it will merge args_groups to arg0
-        # inference data object
-        if not inplace:
-            for group in arg0_groups:
-                group_data = getattr(arg0, group)
-                args_groups[group] = deepcopy(group_data) if copy else group_data
-
-        other_groups = [group for group in args_groups if group not in SUPPORTED_GROUPS_ALL]
-
-        for group in SUPPORTED_GROUPS_ALL + other_groups:
-            if group not in args_groups:
-                continue
-            if inplace:
-                if group.startswith(WARMUP_TAG):
-                    arg0._groups_warmup.append(group)
-                else:
-                    arg0._groups.append(group)
-                setattr(arg0, group, args_groups[group])
-            else:
-                inference_data_dict[group] = args_groups[group]
-        if inplace:
-            other_groups = [
-                group for group in arg0_groups if group not in SUPPORTED_GROUPS_ALL
-            ] + other_groups
-            sorted_groups = [
-                group for group in SUPPORTED_GROUPS + other_groups if group in arg0._groups
-            ]
-            setattr(arg0, "_groups", sorted_groups)
-            sorted_groups_warmup = [
-                group
-                for group in SUPPORTED_GROUPS_WARMUP + other_groups
-                if group in arg0._groups_warmup
-            ]
-            setattr(arg0, "_groups_warmup", sorted_groups_warmup)
-    else:
-        arg0 = args[0]
-        arg0_groups = arg0._groups_all
-        for arg in args[1:]:
-            for group0 in arg0_groups:
-                if group0 not in arg._groups_all:
-                    if group0 == "observed_data":
-                        continue
-                    msg = "Mismatch between the groups."
-                    raise TypeError(msg)
-            for group in arg._groups_all:
-                # handle data groups separately
-                if group not in ["observed_data", "constant_data", "predictions_constant_data"]:
-                    # assert that groups are equal
-                    if group not in arg0_groups:
-                        msg = "Mismatch between the groups."
-                        raise TypeError(msg)
-
-                    # assert that variables are equal
-                    group_data = getattr(arg, group)
-                    group_vars = group_data.data_vars
-
-                    if not inplace and group in inference_data_dict:
-                        group0_data = inference_data_dict[group]
-                    else:
-                        group0_data = getattr(arg0, group)
-                    group0_vars = group0_data.data_vars
-
-                    for var in group0_vars:
-                        if var not in group_vars:
-                            msg = "Mismatch between the variables."
-                            raise TypeError(msg)
-
-                    for var in group_vars:
-                        if var not in group0_vars:
-                            msg = "Mismatch between the variables."
-                            raise TypeError(msg)
-                        var_dims = group_data[var].dims
-                        var0_dims = group0_data[var].dims
-                        if var_dims != var0_dims:
-                            msg = "Mismatch between the dimensions."
-                            raise TypeError(msg)
-
-                        if dim not in var_dims or dim not in var0_dims:
-                            msg = f"Dimension {dim} missing."
-                            raise TypeError(msg)
-
-                    # xr.concat
-                    concatenated_group = xr.concat((group0_data, group_data), dim=dim)
-                    if reset_dim:
-                        concatenated_group[dim] = range(concatenated_group[dim].size)
-
-                    # handle attrs
-                    if hasattr(group0_data, "attrs"):
-                        group0_attrs = deepcopy(getattr(group0_data, "attrs"))
-                    else:
-                        group0_attrs = OrderedDict()
-
-                    if hasattr(group_data, "attrs"):
-                        group_attrs = getattr(group_data, "attrs")
-                    else:
-                        group_attrs = {}
-
-                    # gather attrs results to group0_attrs
-                    for attr_key, attr_values in group_attrs.items():
-                        group0_attr_values = group0_attrs.get(attr_key, None)
-                        equality = attr_values == group0_attr_values
-                        if hasattr(equality, "__iter__"):
-                            equality = np.all(equality)
-                        if equality:
-                            continue
-                        # handle special cases:
-                        if attr_key in ("created_at", "previous_created_at"):
-                            # check the defaults
-                            if not hasattr(group0_attrs, "previous_created_at"):
-                                group0_attrs["previous_created_at"] = []
-                                if group0_attr_values is not None:
-                                    group0_attrs["previous_created_at"].append(group0_attr_values)
-                            # check previous values
-                            if attr_key == "previous_created_at":
-                                if not isinstance(attr_values, list):
-                                    attr_values = [attr_values]
-                                group0_attrs["previous_created_at"].extend(attr_values)
-                                continue
-                            # update "created_at"
-                            if group0_attr_values != current_time:
-                                group0_attrs[attr_key] = current_time
-                            group0_attrs["previous_created_at"].append(attr_values)
-
-                        elif attr_key in group0_attrs:
-                            combined_key = f"combined_{attr_key}"
-                            if combined_key not in group0_attrs:
-                                group0_attrs[combined_key] = [group0_attr_values]
-                            group0_attrs[combined_key].append(attr_values)
-                        else:
-                            group0_attrs[attr_key] = attr_values
-                    # update attrs
-                    setattr(concatenated_group, "attrs", group0_attrs)
-
-                    if inplace:
-                        setattr(arg0, group, concatenated_group)
-                    else:
-                        inference_data_dict[group] = concatenated_group
-                else:
-                    # observed_data, "constant_data", "predictions_constant_data",
-                    if group not in arg0_groups:
-                        setattr(arg0, group, deepcopy(group_data) if copy else group_data)
-                        arg0._groups.append(group)
-                        continue
-
-                    # assert that variables are equal
-                    group_data = getattr(arg, group)
-                    group_vars = group_data.data_vars
-
-                    group0_data = getattr(arg0, group)
-                    if not inplace:
-                        group0_data = deepcopy(group0_data)
-                    group0_vars = group0_data.data_vars
-
-                    for var in group_vars:
-                        if var not in group0_vars:
-                            var_data = group_data[var]
-                            getattr(arg0, group)[var] = var_data
-                        else:
-                            var_data = group_data[var]
-                            var0_data = group0_data[var]
-                            if dim in var_data.dims and dim in var0_data.dims:
-                                concatenated_var = xr.concat((group_data, group0_data), dim=dim)
-                                group0_data[var] = concatenated_var
-
-                    # handle attrs
-                    if hasattr(group0_data, "attrs"):
-                        group0_attrs = getattr(group0_data, "attrs")
-                    else:
-                        group0_attrs = OrderedDict()
-
-                    if hasattr(group_data, "attrs"):
-                        group_attrs = getattr(group_data, "attrs")
-                    else:
-                        group_attrs = {}
-
-                    # gather attrs results to group0_attrs
-                    for attr_key, attr_values in group_attrs.items():
-                        group0_attr_values = group0_attrs.get(attr_key, None)
-                        equality = attr_values == group0_attr_values
-                        if hasattr(equality, "__iter__"):
-                            equality = np.all(equality)
-                        if equality:
-                            continue
-                        # handle special cases:
-                        if attr_key in ("created_at", "previous_created_at"):
-                            # check the defaults
-                            if not hasattr(group0_attrs, "previous_created_at"):
-                                group0_attrs["previous_created_at"] = []
-                                if group0_attr_values is not None:
-                                    group0_attrs["previous_created_at"].append(group0_attr_values)
-                            # check previous values
-                            if attr_key == "previous_created_at":
-                                if not isinstance(attr_values, list):
-                                    attr_values = [attr_values]
-                                group0_attrs["previous_created_at"].extend(attr_values)
-                                continue
-                            # update "created_at"
-                            if group0_attr_values != current_time:
-                                group0_attrs[attr_key] = current_time
-                            group0_attrs["previous_created_at"].append(attr_values)
-
-                        elif attr_key in group0_attrs:
-                            combined_key = f"combined_{attr_key}"
-                            if combined_key not in group0_attrs:
-                                group0_attrs[combined_key] = [group0_attr_values]
-                            group0_attrs[combined_key].append(attr_values)
-
-                        else:
-                            group0_attrs[attr_key] = attr_values
-                    # update attrs
-                    setattr(group0_data, "attrs", group0_attrs)
-
-                    if inplace:
-                        setattr(arg0, group, group0_data)
-                    else:
-                        inference_data_dict[group] = group0_data
-
-    if not inplace:
-        inference_data_dict["attrs"] = combined_attr
-
-    return None if inplace else InferenceData(**inference_data_dict)