Preserve observed_data coordinates in weight_predictions

src/arviz_stats/manipulation.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import xarray as xr
-from arviz_base import convert_to_dataset, convert_to_datatree, extract, from_dict, rcParams
+from arviz_base import convert_to_dataset, convert_to_datatree, extract, rcParams
 
 from arviz_stats.utils import get_array_function
 from arviz_stats.validate import validate_dims
@@ -193,12 +193,10 @@ def weight_predictions(
         for samples, dt in zip(new_samples, dts)
     ]
 
-    weighted_samples = from_dict(
+    weighted_samples = xr.DataTree.from_dict(
         {
             group: xr.concat(new_idatas, dim="sample"),
             "observed_data": dts[0].observed_data,
-        },
-        sample_dims=list(new_idatas[0].dims.keys()),
+        }
     )
-
     return weighted_samples

tests/test_manipulation.py

@@ -6,10 +6,11 @@
 
 from .helpers import importorskip
 
-azb = importorskip("arviz_base")
 xr = importorskip("xarray")
+azb = importorskip("arviz_base")
+
 
-from arviz_stats.manipulation import thin
+from arviz_stats.manipulation import thin, weight_predictions
 
 
 @pytest.mark.parametrize("factor", [2, 5, 10])
@@ -120,7 +121,21 @@ def test_thin_invalid_factor(datatree):
         thin(data, factor=0)
 
 
-def test_thin_negative_factor(datatree):
-    data = datatree.posterior["mu"]
-    with pytest.raises(ValueError, match="factor must be greater than 1"):
-        thin(data, factor=-5)
+def test_weight_predictions_preserves_coords():
+    dt = azb.load_arviz_data("centered_eight")
+
+    # Add NON-index coordinate to all groups
+    dt = dt.map_over_datasets(
+        lambda ds: ds.assign_coords(school=("school", ["a", "b", "c", "d", "e", "f", "g", "h"]))
+    )
+
+    out = weight_predictions([dt, dt])
+
+    # ---- observed_data ----
+    for coord in dt.observed_data.coords:
+        assert coord in out.observed_data.coords
+        assert (out.observed_data.coords[coord] == dt.observed_data.coords[coord]).all()
+
+    # ---- posterior_predictive ----
+    for coord in dt.posterior_predictive.coords:
+        assert coord in out.posterior_predictive.coords