Fixed UserWarning when converting sample_stats to idata

[MSK-005]

Description

This PR aims to close #7821. Previously, in some cases, running pm.sample_smc() resulted in a UserWarning being generated:

UserWarning: More chains (5) than draws (1). Passed array should have shape (chains, draws, *shape) warnings.warn(

What I did was update the dict_to_dataset() function call(s) in pymc/backends/arviz.py and pymc/smc/sampling.py. In pymc/backends/arviz.py, I changed the import statement so that dict_to_dataset() was imported from arviz_base rather than arviz.data.base, and the rest of the changes just involved updating the function call. Another important detail is that in pymc/smc/sampling/py, I added another argument when calling dict_to_dataset, which is sample_dims=["chain"], since the variable sample_stats_dict was a 1-D dict.

Related Issue

• Closes #7821
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[MSK-005]

Can anyone explain how to fix these checks? pre-commit passes all, but here, one test is failing.

[williambdean]

Both of the requirements.txt files should not be manually edited (see comment at top of the file). If you need to add dependencies, change conda-envs/environment-dev.yml

[ricardoV94]

I'm not sure what's the status of arviz-base and if it's stable enough to switch as a dependency. This would probably require a major PyMC version bump as well?

CC @OriolAbril and @aloctavodia

[aloctavodia]

In general we could begin a migration to "New ArviZ", but to reduce friction, we should do it in a coordinated manner and change arviz imports everywhere, not just here. It may be a good idea to have a dedicated PyMC release to work on the transition.

In particular, about this PR, from_dict returns a DataTree object, not an InferenceData object. DataTree is a native xarray object that effectively replaces InferenceData in ArviZ. Many functions (all that make sense?) in ArviZ still understand InferenceData objects, but only for backward compatibility.

[OriolAbril]

I will take a closer look in the coming days but it looks good. We can have both arviz and arviz-base installed together without problem and I think doing that doesn't require a major version because for the most part, users won't really notice any change.

In this case it is only a switch between the two dict_to_dataset whose only difference are a couple function arguments, that users can't change anyway.

The main things to comment are trying to switch as many imports as possible to arviz-base instead of combining arviz-base and arviz.data imports in the same file. And requiring arviz-base newer than 0.7

[MSK-005]

I'll wait for Mr Oriol's comments in the coming days. If you guys feel like it's too early to shift to arviz-base, I can revert the changes and wrap the function call in an ignore-warnings context. Thank you, all of you, for your time.