Skip to content

Comments

Preserve observed_data coordinates in weight_predictions#285

Open
ShivamRajSri wants to merge 2 commits intoarviz-devs:mainfrom
ShivamRajSri:fix-weight-predictions-observed-coords
Open

Preserve observed_data coordinates in weight_predictions#285
ShivamRajSri wants to merge 2 commits intoarviz-devs:mainfrom
ShivamRajSri:fix-weight-predictions-observed-coords

Conversation

@ShivamRajSri
Copy link

@ShivamRajSri ShivamRajSri commented Jan 15, 2026

Closes #269

  • Preserve observe_data coordinates
  • Add test

@read-the-docs-community
Copy link

Documentation build overview

📚 arviz-stats | 🛠️ Build #31021022 | 📁 Comparing e9676bd against latest (b2de789)


🔍 Preview build

Show files changed (8 files in total): 📝 8 modified | ➕ 0 added | ➖ 0 deleted
File Status
_modules/arviz_stats/manipulation.html 📝 modified
api/generated/arviz_stats.eti.html 📝 modified
api/generated/arviz_stats.hdi.html 📝 modified
api/generated/arviz_stats.histogram.html 📝 modified
api/generated/arviz_stats.kde.html 📝 modified
api/generated/arviz_stats.loo_kfold.html 📝 modified
api/generated/arviz_stats.mode.html 📝 modified
api/generated/arviz_stats.qds.html 📝 modified

@ShivamRajSri
Copy link
Author

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!

@aloctavodia aloctavodia changed the title FIX: preserve observed_data coordinates in weight_predictions Preserve observed_data coordinates in weight_predictions Feb 2, 2026
sample_dims=list(new_idatas[0].dims.keys()),
)

weighted_samples["observed_data"] = dts[0].observed_data
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally the test should also add non-indexing coordinates to both posterior predictive and observed data to check these are not lost either

Comment on lines +141 to +143
assert set(out.observed_data.coords) == set(original_coords)
for coord in original_coords:
assert (out.observed_data.coords[coord] == original_coords[coord]).all()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checks need to happen for both observed_data and posterior_predictive group

Copy link
Member

@OriolAbril OriolAbril left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sorry for the slow review, let me know if there are any doubts about the comments

@ShivamRajSri
Copy link
Author

Thanks for the detailed review @OriolAbril,that clarifies the issue well.
I see the two sources of coordinate loss now: the use of .reset_coords() during extraction and reconstructing labeled xarray objects via arviz_base.from_dict.
I’ll update the implementation to avoid from_dict here and instead build the output with xarray.DataTree.from_dict, so dims and non-index coordinates are preserved correctly.
I’ll also extend the test to add non-index coordinates and check preservation for both observed_data and posterior_predictive, and push an updated commit shortly.

@ShivamRajSri
Copy link
Author

@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.
I also extended the tests to explicitly add non-index coordinates and assert that both observed_data and posterior_predictive coordinates are preserved after calling weight_predictions.
I’ve pushed the updated commit — please let me know if this matches the intended behavior or if you’d like the test coverage adjusted further.

@ShivamRajSri
Copy link
Author

RTD docs build failed; I’m looking into the Sphinx warning causing it (likely related to weight_predictions docstring / imports).
Will push a follow-up shortly once fixed.

@ShivamRajSri
Copy link
Author

The RTD failure is due to an incompatibility between Sphinx 9.1.0 and sphinx_autosummary_accessors
This is unrelated to the changes in this PR. Happy to help pin Sphinx or adjust the docs setup if you’d like to handle it here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Double check weight_predictions implementation

2 participants