Repository: arviz-devs/arviz

Issue:
add merge/add group function
## Tell us about it
`az.concat` currently raises an error if there are repeated groups which is perfectly fine, however, there are situations where one wants to add the different groups to the original idata object and discard the repeated groups. 

related to #1066 

## Thoughts on implementation

several options, some of them complementary:

1. add an `idata.merge` function so that idata wiht groups post, sample_stats, obs_data can be merged with idata with groups post_pred and obs_data discarding the obs_data in one of the two (could have left/right modes)
2. add an `idata.add_groups` so that we can pass it dicts (calls dict_to_dataset) or datasets as kwargs and they are added to idata (I think kwargs is important to allow special usecases and support any group_name, alternatively a mapping {group_name: ds/dict} could be used)
3. make `idata.concat` check equality of datasets with isclose on both vals and coords (but not on attrs! they will surely have different creation time)



Discussion:
tagging @percygautam 

Resolving patch:
diff --git a/arviz/data/inference_data.py b/arviz/data/inference_data.py
--- a/arviz/data/inference_data.py
+++ b/arviz/data/inference_data.py
@@ -12,7 +12,9 @@
 import numpy as np
 import xarray as xr
 from xarray.core.options import OPTIONS
+from xarray.core.utils import either_dict_or_kwargs
 
+from .base import dict_to_dataset
 from .base import _extend_xr_method
 from ..utils import _subset_list, HtmlTemplate
 from ..rcparams import rcParams
@@ -627,6 +629,101 @@ def rename_dims(self, name_dict=None, groups=None, filter_groups=None, inplace=F
         else:
             return out
 
+    def add_groups(self, group_dict=None, coords=None, dims=None, **kwargs):
+        """Add new groups to InferenceData object.
+
+        Parameters
+        ----------
+        group_dict: dict of {str : dict or xarray.Dataset}, optional
+            Groups to be added
+        coords : dict[str] -> ndarray
+            Coordinates for the dataset
+        dims : dict[str] -> list[str]
+            Dimensions of each variable. The keys are variable names, values are lists of
+            coordinates.
+        **kwargs: mapping
+            The keyword arguments form of group_dict. One of group_dict or kwargs must be provided.
+
+        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("{} group(s) already exists.".format(repeated_groups))
+        for group, dataset in group_dict.items():
+            if group not in SUPPORTED_GROUPS_ALL:
+                warnings.warn(
+                    "The group {} is not defined in the InferenceData scheme".format(group),
+                    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):
+                    self._groups_warmup.append(group)
+                else:
+                    self._groups.append(group)
+
+    def extend(self, other, join="left"):
+        """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``.
+
+        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("join must be either 'left' or 'right', found {}".format(join))
+        for group in other._groups_all:  # pylint: disable=protected-access
+            if hasattr(self, group):
+                if join == "left":
+                    continue
+            if group not in SUPPORTED_GROUPS_ALL:
+                warnings.warn(
+                    "{} group is not defined in the InferenceData scheme".format(group), UserWarning
+                )
+            dataset = getattr(other, group)
+            setattr(self, group, dataset)
+
     set_index = _extend_xr_method(xr.Dataset.set_index)
     get_index = _extend_xr_method(xr.Dataset.get_index)
     reset_index = _extend_xr_method(xr.Dataset.reset_index)
@@ -886,6 +983,11 @@ def concat(*args, dim=None, copy=True, inplace=False, reset_dim=True):
         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
@@ -977,8 +1079,9 @@ def concat(*args, dim=None, copy=True, inplace=False, reset_dim=True):
                 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"
                     )
-                    msg += " Valid dimensions are `chain` and `draw`."
                     raise TypeError(msg)
                 group_data = getattr(arg, group)
                 args_groups[group] = deepcopy(group_data) if copy else group_data


Test patch:
diff --git a/arviz/tests/base_tests/test_data.py b/arviz/tests/base_tests/test_data.py
--- a/arviz/tests/base_tests/test_data.py
+++ b/arviz/tests/base_tests/test_data.py
@@ -584,6 +584,63 @@ def test_repr_html(self):
         assert escape(idata.__repr__()) in html
         xr.set_options(display_style=display_style)
 
+    def test_add_groups(self, data_random):
+        data = np.random.normal(size=(4, 500, 8))
+        idata = data_random
+        idata.add_groups({"prior": {"a": data[..., 0], "b": data}})
+        assert "prior" in idata._groups  # pylint: disable=protected-access
+        assert isinstance(idata.prior, xr.Dataset)
+        assert hasattr(idata, "prior")
+
+        idata.add_groups(posterior_warmup={"a": data[..., 0], "b": data})
+        assert "posterior_warmup" in idata._groups  # pylint: disable=protected-access
+        assert isinstance(idata.posterior_warmup, xr.Dataset)
+        assert hasattr(idata, "posterior_warmup")
+
+    def test_add_groups_warning(self, data_random):
+        data = np.random.normal(size=(4, 500, 8))
+        idata = data_random
+        with pytest.warns(UserWarning, match="The group.+not defined in the InferenceData scheme"):
+            idata.add_groups({"new_group": idata.posterior})
+        with pytest.warns(UserWarning, match="the default dims.+will be added automatically"):
+            idata.add_groups(constant_data={"a": data[..., 0], "b": data})
+        assert idata.new_group.equals(idata.posterior)
+
+    def test_add_groups_error(self, data_random):
+        idata = data_random
+        with pytest.raises(ValueError, match="One of.+must be provided."):
+            idata.add_groups()
+        with pytest.raises(ValueError, match="Arguments.+xr.Dataset, xr.Dataarray or dicts"):
+            idata.add_groups({"new_group": "new_group"})
+        with pytest.raises(ValueError, match="group.+already exists"):
+            idata.add_groups({"posterior": idata.posterior})
+
+    def test_extend(self, data_random):
+        idata = data_random
+        idata2 = create_data_random(groups=["prior", "prior_predictive", "observed_data"], seed=7)
+        idata.extend(idata2)
+        assert hasattr(idata, "prior")
+        assert hasattr(idata, "prior_predictive")
+        assert idata.prior.equals(idata2.prior)
+        assert not idata.observed_data.equals(idata2.observed_data)
+        assert idata.prior_predictive.equals(idata2.prior_predictive)
+
+        idata.extend(idata2, join="right")
+        assert idata.prior.equals(idata2.prior)
+        assert idata.observed_data.equals(idata2.observed_data)
+        assert idata.prior_predictive.equals(idata2.prior_predictive)
+
+    def test_extend_errors_warnings(self, data_random):
+        idata = data_random
+        idata2 = create_data_random(groups=["prior", "prior_predictive", "observed_data"], seed=7)
+        with pytest.raises(ValueError, match="Extending.+InferenceData objects only."):
+            idata.extend("something")
+        with pytest.raises(ValueError, match="join must be either"):
+            idata.extend(idata2, join="outer")
+        idata2.add_groups(new_group=idata2.prior)
+        with pytest.warns(UserWarning):
+            idata.extend(idata2)
+
 
 class TestNumpyToDataArray:
     def test_1d_dataset(self):