Allow xarray.Dataarray input to plots.

CHANGELOG.md

@@ -7,6 +7,7 @@
 * Add out-of-sample groups (`predictions` and `predictions_constant_data`) and `constant_data` group to pyro translation #1090
 * Add `num_chains` and `pred_dims` arguments to io_pyro #1090
 * Integrate jointplot into pairplot, add point-estimate and overlay of plot kinds (#1079)
+* Allow xarray.Dataarray input for plots.(#1120)
 ### Maintenance and fixes
 * Fixed behaviour of `credible_interval=None` in `plot_posterior` (#1115)
 * Fixed hist kind of `plot_dist` with multidimensional input (#1115)

arviz/data/converters.py

@@ -33,6 +33,8 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
             | emcee sampler: Automatically extracts data
             | pyro MCMC: Automatically extracts data
             | xarray.Dataset: adds to InferenceData as only group
+            | xarray.DataArray: creates an xarray dataset as the only group, gives the
+                         array an arbitrary name, if name not set
             | dict: creates an xarray dataset as the only group
             | numpy array: creates an xarray dataset as the only group, gives the
                          array an arbitrary name
@@ -97,6 +99,10 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
     # Cases that convert to xarray
     if isinstance(obj, xr.Dataset):
         dataset = obj
+    elif isinstance(obj, xr.DataArray):
+        if obj.name is None:
+            obj.name = 'x'
+        dataset = obj.to_dataset()
     elif isinstance(obj, dict):
         dataset = dict_to_dataset(obj, coords=coords, dims=dims)
     elif isinstance(obj, np.ndarray):
@@ -109,6 +115,7 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
         return from_cmdstan(**kwargs)
     else:
         allowable_types = (
+            "xarray dataarray",
             "xarray dataset",
             "dict",
             "netcdf filename",
@@ -149,6 +156,8 @@ def convert_to_dataset(obj, *, group="posterior", coords=None, dims=None):
             pystan fit: Automatically extracts data
             pymc3 trace: Automatically extracts data
             xarray.Dataset: adds to InferenceData as only group
+            xarray.DataArray: creates an xarray dataset as the only group, gives the
+                         array an arbitrary name, if name not set
             dict: creates an xarray dataset as the only group
             numpy array: creates an xarray dataset as the only group, gives the
                          array an arbitrary name

arviz/tests/base_tests/test_data.py

@@ -873,3 +873,41 @@ def test_id_conversion_args(self):
         assert isinstance(inference_data, InferenceData)
         assert set(inference_data.posterior.coords["Ivies"].values) == set(IVIES)
         assert inference_data.posterior["theta"].dims == ("chain", "draw", "Ivies")
+
+
+
+class TestDataArrayToDataset:
+    def test_1d_dataset(self):
+        size = 100
+        dataset = convert_to_dataset(xr.DataArray(np.random.randn(1, size), dims=('chain', 'draw')))
+        assert len(dataset.data_vars) == 1
+
+        assert dataset.chain.shape == (1, )
+        assert dataset.draw.shape == (size, )
+
+
+    def test_nd_to_dataset(self):
+        shape = (1, 2, 3, 4, 5)
+        dataset = convert_to_dataset(
+            xr.DataArray(np.random.randn(*shape),
+            dims=('chain', 'draw', 'dim_0', 'dim_1', 'dim_2')))
+        assert len(dataset.data_vars) == 1
+        var_name = list(dataset.data_vars)[0]
+
+        assert dataset.chain.shape == shape[:1]
+        assert dataset.draw.shape == shape[1:2]
+        assert dataset[var_name].shape == shape
+
+    def test_nd_to_inference_data(self):
+        shape = (1, 2, 3, 4, 5)
+        inference_data = convert_to_inference_data(
+            xr.DataArray(np.random.randn(*shape),
+            dims=('chain', 'draw', 'dim_0', 'dim_1', 'dim_2')), group="prior")
+        assert hasattr(inference_data, "prior")
+        assert len(inference_data.prior.data_vars) == 1
+        var_name = list(inference_data.prior.data_vars)[0]
+
+        assert inference_data.prior.chain.shape == shape[:1]
+        assert inference_data.prior.draw.shape == shape[1:2]
+        assert inference_data.prior[var_name].shape == shape
+        assert repr(inference_data).startswith("Inference data with groups")