Preserve observed_data coordinates in weight_predictions

[ShivamRajSri]

Closes #269

• Preserve observe_data coordinates
• Add test

[read-the-docs-community]

Documentation build overview

📚 arviz-stats | 🛠️ Build #31815498 | 📁 Comparing 32b2e4f against latest (0738be9)

🔍 Preview build

Show files changed (13 files in total): 📝 13 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.iqr.html	📝 modified
api/generated/arviz_stats.kde.html	📝 modified
api/generated/arviz_stats.loo_kfold.html	📝 modified
api/generated/arviz_stats.mad.html	📝 modified
api/generated/arviz_stats.mean.html	📝 modified
api/generated/arviz_stats.median.html	📝 modified
api/generated/arviz_stats.mode.html	📝 modified
api/generated/arviz_stats.qds.html	📝 modified
api/generated/arviz_stats.var.html	📝 modified

[ShivamRajSri]

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!

[OriolAbril]

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

[ShivamRajSri]

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]

@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]

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]

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.

[OriolAbril]

sorry about the slow review cadence. We have just released the 1.0 version so hopefully things will improve now, but probably not immediately, we do have a significant backlog to catch up with before that.

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 have added a comment mentioning that, but this is not the case at all.

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.

We already included the pin. If you rebase on main the RTD build will go back to green again.

[OriolAbril]

Suggested change

dt = convert_to_datatree(idata)

idata variable is already a DataTree. We often use idata or dt but the type is always DataTree. The use of idata is mostly for historical reasons, but there is also no need to change that because a DataTree object can have an arbitrary schema regarding groups and variables while our variables and inputs follow a stricter set of conventions (defined here)

[OriolAbril]

I probably wasn't very clear so I'll try again. idata is already a DataTree so calling convert_to_datatree is useless and should not be done (hence the suggestion within the review comment to remove the line). The rest of the comment is because it doesn't matter if you keep idata = az.load_arviz_data then use idata later on or update to dt = az.load_arviz_data.

[OriolAbril]

there are still comments to address. I'd recommend going over all of them again, note you can mark them as resolved so it is easier to see what is still pending.

And all CI jobs are failing not sure if only due to but at least in part due to the imports in the test file. Look at what the file is already doing and at https://python.arviz.org/projects/stats/en/latest/contributing/testing.html

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.85%. Comparing base (0738be9) to head (32b2e4f).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #285      +/-   ##
==========================================
+ Coverage   84.59%   84.85%   +0.25%     
==========================================
  Files          43       43              
  Lines        5974     5974              
==========================================
+ Hits         5054     5069      +15     
+ Misses        920      905      -15     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ShivamRajSri]

Hi! @OriolAbril I’ve applied all the requested changes and the tests and lint checks are passing now.

The only remaining failure is the ReadTheDocs build. From the logs it looks like a Sphinx compatibility issue with sphinx_autosummary_accessors, which seems unrelated to the changes in this PR.

Please let me know if there’s something I should modify on my side. Thanks!

[OriolAbril]

I have unresolved this comment because it was not addressed.

Extra note: if you can please rebase on main. As I mentioned in another comment this will fix the RTD build.

[ShivamRajSri]

Hi @OriolAbril,

Sorry for missing that comment earlier. I’ve been in the middle of my mid-semester exams recently, so I wasn’t able to go through all the review comments as carefully as I should have. That’s my mistake.

Thanks for pointing it out again. I’ll address the unresolved comment and rebase the branch on main as suggested so the RTD build should also be fixed.

Appreciate your patience!

[ShivamRajSri]

Thanks @OriolAbril for pointing that out. I’ve updated the test to remove the unnecessary convert_to_datatree call and rebased the branch on main as suggested.

Happy to make further changes if needed.