edit groups in idata.extend

CHANGELOG.md

@@ -8,6 +8,8 @@
 * Added `circ_var_names` argument to `plot_trace` allowing for circular traceplot (Matplotlib) ([1336](https://github.com/arviz-devs/arviz/pull/1336))
 * Ridgeplot is hdi aware. By default displays truncated densities at the specified `hdi_prop` level ([1348](https://github.com/arviz-devs/arviz/pull/1348))
 * Added `plot_separation` ([1359](https://github.com/arviz-devs/arviz/pull/1359))
+* Extended methods from `xr.Dataset` to `InferenceData` ([1254](https://github.com/arviz-devs/arviz/pull/1254))
+* Add `extend` and `add_groups` to `InferenceData` ([1300](https://github.com/arviz-devs/arviz/pull/1300) and [1386](https://github.com/arviz-devs/arviz/pull/1386))
 
 
 ### Maintenance and fixes
@@ -27,6 +29,8 @@
 * Add more `plot_parallel` examples ([#1380](https://github.com/arviz-devs/arviz/pull/1380))
 * Bump minimum xarray version to 0.16.1 ([#1389](https://github.com/arviz-devs/arviz/pull/1389)
 * Fix multi rope for `plot_forest` ([#1390](https://github.com/arviz-devs/arviz/pull/1390))
+* Bump minimum xarray version to 0.16.1 ([#1389](https://github.com/arviz-devs/arviz/pull/1389))
+* `from_dict` will now store warmup groups even with the main group missing ([1386](https://github.com/arviz-devs/arviz/pull/1386))
 
 ### Deprecation
 

arviz/data/base.py

@@ -245,6 +245,7 @@ def make_attrs(attrs=None, library=None):
 
 def _extend_xr_method(func):
     """Make wrapper to extend methods from xr.Dataset to InferenceData Class."""
+    # pydocstyle requires a non empty line
 
     @functools.wraps(func)
     def wrapped(self, *args, **kwargs):
@@ -280,7 +281,7 @@ def wrapped(self, *args, **kwargs):
         metagroup names. A la `pandas.filter`.
     inplace: bool, optional
         If ``True``, modify the InferenceData object inplace,
-        otherwise, return the modified copy. 
+        otherwise, return the modified copy.
     """
     see_also = """
     See Also

arviz/data/inference_data.py

@@ -820,6 +820,10 @@ def extend(self, other, join="left"):
                 )
             dataset = getattr(other, group)
             setattr(self, group, dataset)
+            if group.startswith(WARMUP_TAG):
+                self._groups_warmup.append(group)
+            else:
+                self._groups.append(group)
 
     set_index = _extend_xr_method(xr.Dataset.set_index)
     get_index = _extend_xr_method(xr.Dataset.get_index)

arviz/data/io_dict.py

@@ -6,7 +6,7 @@
 from .. import utils
 from ..rcparams import rcParams
 from .base import dict_to_dataset, generate_dims_coords, make_attrs, requires
-from .inference_data import InferenceData
+from .inference_data import InferenceData, WARMUP_TAG
 
 
 # pylint: disable=too-many-instance-attributes
@@ -69,11 +69,15 @@ def __init__(
         self.attrs.pop("created_at", None)
         self.attrs.pop("arviz_version", None)
 
-    @requires("posterior")
+    def _init_dict(self, attr_name):
+        dict_or_none = getattr(self, attr_name, None)
+        return {} if dict_or_none is None else dict_or_none
+
+    @requires(["posterior", f"{WARMUP_TAG}posterior"])
     def posterior_to_xarray(self):
         """Convert posterior samples to xarray."""
-        data = self.posterior
-        data_warmup = self.warmup_posterior if self.warmup_posterior is not None else {}
+        data = self._init_dict("posterior")
+        data_warmup = self._init_dict(f"{WARMUP_TAG}posterior")
         if not isinstance(data, dict):
             raise TypeError("DictConverter.posterior is not a dictionary")
         if not isinstance(data_warmup, dict):
@@ -95,11 +99,11 @@ def posterior_to_xarray(self):
             ),
         )
 
-    @requires("sample_stats")
+    @requires(["sample_stats", f"{WARMUP_TAG}sample_stats"])
     def sample_stats_to_xarray(self):
         """Convert sample_stats samples to xarray."""
-        data = self.sample_stats
-        data_warmup = self.warmup_sample_stats if self.warmup_sample_stats is not None else {}
+        data = self._init_dict("sample_stats")
+        data_warmup = self._init_dict(f"{WARMUP_TAG}sample_stats")
         if not isinstance(data, dict):
             raise TypeError("DictConverter.sample_stats is not a dictionary")
         if not isinstance(data_warmup, dict):
@@ -122,11 +126,11 @@ def sample_stats_to_xarray(self):
             ),
         )
 
-    @requires("log_likelihood")
+    @requires(["log_likelihood", f"{WARMUP_TAG}log_likelihood"])
     def log_likelihood_to_xarray(self):
         """Convert log_likelihood samples to xarray."""
-        data = self.log_likelihood
-        data_warmup = self.warmup_log_likelihood if self.warmup_log_likelihood is not None else {}
+        data = self._init_dict("log_likelihood")
+        data_warmup = self._init_dict(f"{WARMUP_TAG}log_likelihood")
         if not isinstance(data, dict):
             raise TypeError("DictConverter.log_likelihood is not a dictionary")
         if not isinstance(data_warmup, dict):
@@ -141,13 +145,11 @@ def log_likelihood_to_xarray(self):
             ),
         )
 
-    @requires("posterior_predictive")
+    @requires(["posterior_predictive", f"{WARMUP_TAG}posterior_predictive"])
     def posterior_predictive_to_xarray(self):
         """Convert posterior_predictive samples to xarray."""
-        data = self.posterior_predictive
-        data_warmup = (
-            self.warmup_posterior_predictive if self.warmup_posterior_predictive is not None else {}
-        )
+        data = self._init_dict("posterior_predictive")
+        data_warmup = self._init_dict(f"{WARMUP_TAG}posterior_predictive")
         if not isinstance(data, dict):
             raise TypeError("DictConverter.posterior_predictive is not a dictionary")
         if not isinstance(data_warmup, dict):
@@ -162,11 +164,11 @@ def posterior_predictive_to_xarray(self):
             ),
         )
 
-    @requires("predictions")
+    @requires(["predictions", f"{WARMUP_TAG}predictions"])
     def predictions_to_xarray(self):
         """Convert predictions to xarray."""
-        data = self.predictions
-        data_warmup = self.warmup_predictions if self.warmup_predictions is not None else {}
+        data = self._init_dict("predictions")
+        data_warmup = self._init_dict(f"{WARMUP_TAG}predictions")
         if not isinstance(data, dict):
             raise TypeError("DictConverter.predictions is not a dictionary")
         if not isinstance(data_warmup, dict):

arviz/tests/base_tests/test_data.py

@@ -624,10 +624,10 @@ def test_add_groups(self, data_random):
         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")
+        idata.add_groups(warmup_posterior={"a": data[..., 0], "b": data})
+        assert "warmup_posterior" in idata._groups_all  # pylint: disable=protected-access
+        assert isinstance(idata.warmup_posterior, xr.Dataset)
+        assert hasattr(idata, "warmup_posterior")
 
     def test_add_groups_warning(self, data_random):
         data = np.random.normal(size=(4, 500, 8))
@@ -649,8 +649,12 @@ def test_add_groups_error(self, data_random):
 
     def test_extend(self, data_random):
         idata = data_random
-        idata2 = create_data_random(groups=["prior", "prior_predictive", "observed_data"], seed=7)
+        idata2 = create_data_random(
+            groups=["prior", "prior_predictive", "observed_data", "warmup_posterior"], seed=7
+        )
         idata.extend(idata2)
+        assert "prior" in idata._groups_all  # pylint: disable=protected-access
+        assert "warmup_posterior" in idata._groups_all  # pylint: disable=protected-access
         assert hasattr(idata, "prior")
         assert hasattr(idata, "prior_predictive")
         assert idata.prior.equals(idata2.prior)

arviz/tests/helpers.py

@@ -157,6 +157,7 @@ def create_data_random(groups=None, seed=10):
         prior_predictive={"a": data[..., 0], "b": data},
         warmup_posterior={"a": data[..., 0], "b": data},
         warmup_posterior_predictive={"a": data[..., 0], "b": data},
+        warmup_prior={"a": data[..., 0], "b": data},
     )
     idata = from_dict(
         **{group: ary for group, ary in idata_dict.items() if group in groups}, save_warmup=True