Preserve observed_data coordinates in weight_predictions#285
Preserve observed_data coordinates in weight_predictions#285ShivamRajSri wants to merge 2 commits intoarviz-devs:mainfrom
Conversation
Documentation build overview
Show files changed (8 files in total): 📝 8 modified | ➕ 0 added | ➖ 0 deleted
|
|
Hi @aloctavodia , just a gentle follow-up on this PR. This fixes the issue where weight_predictions was dropping the original observed_data coordinates and replaces them with a dummy index. The change preserves the original coordinates (e.g. school in centered_eight) and adds a small regression test to lock in the behavior. All tests pass locally, and the change is isolated to weight_predictions. Whenever you have time, I’d appreciate a review — happy to adjust anything if needed. Thanks! |
src/arviz_stats/manipulation.py
Outdated
| sample_dims=list(new_idatas[0].dims.keys()), | ||
| ) | ||
|
|
||
| weighted_samples["observed_data"] = dts[0].observed_data |
There was a problem hiding this comment.
seeing this makes me think there are two layers to the loss of coordinates/dimensions. One is the .reset_coords() a few lines above that I mentioned on the issue, the other is using arviz_base.from_dict. This second one is related to arviz-devs/arviz-base#146 and in this case the fix is to use xarray.DataTree.from_dict instead. Our from_dict expects array_like not things that are already labeled, if we already have xarray objects like we do here then there is no point in using arviz_base.from_dict
|
|
||
|
|
||
| def test_weight_predictions_preserves_observed_coords(): | ||
| idata = az.load_arviz_data("centered_eight") |
There was a problem hiding this comment.
Ideally the test should also add non-indexing coordinates to both posterior predictive and observed data to check these are not lost either
| assert set(out.observed_data.coords) == set(original_coords) | ||
| for coord in original_coords: | ||
| assert (out.observed_data.coords[coord] == original_coords[coord]).all() |
There was a problem hiding this comment.
Checks need to happen for both observed_data and posterior_predictive group
OriolAbril
left a comment
There was a problem hiding this comment.
sorry for the slow review, let me know if there are any doubts about the comments
|
Thanks for the detailed review @OriolAbril,that clarifies the issue well. |
|
@OriolAbril ,I’ve updated the implementation to avoid reconstructing labeled objects with arviz_base.from_dict. Instead, the output is now built using xarray.DataTree.from_dict, which preserves both dimensions and non-index coordinates. |
|
RTD docs build failed; I’m looking into the Sphinx warning causing it (likely related to weight_predictions docstring / imports). |
|
The RTD failure is due to an incompatibility between Sphinx 9.1.0 and sphinx_autosummary_accessors |
Closes #269