Collect sampler warnings only through stats

[michaelosthege]

What is this PR about?
Deleting code that gave special treatments to sampler warnings.

This PR mostly changes three things:

• The SamplerWarning class and related code, including the code for run_convergence_checks is moved from report.py to convergence.py. Here it is more functionally structure, and no longer coupled to the state of a SamplerReport.
• A new stat called "warning" is introduced. This is an object stat to which samplers can emit warnings instead of piggy-backing them onto the draws.
• Handling warnings is removed from the BaseTrace, in favor of managing them via the "warning" stat.
• updating the minimum ArviZ requirement to the latest version that can handle object-typed variables (by skipping and warning about them) when saving (related Support sparse sample stats #6194).
• A pm.sample(keep_warning_stat={False, True}) setting was added to enable access to sampler warnings without breaking .to_netcdf() by default.

The logic in sampling.py was modified to stop handling warnings separately.
Instead, the record_and_warn function now takes over the task of logging warnings that are coming through via the "warning" stat.

I also modified the corresponding tests to be more targeted.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• Detailed information about divergent samples is no longer captured.
• Support for handling sampler warnings via ._warnings attributes on samplers and BaseTrace was removed in favor of a "warning" stat.

Bugfixes / New features

• HMC/NUTS samplers now emit a "warning" stat that ends up in the InferenceData.
• With pm.sample(keep_warning_stat=True) one can now access detailed info about divergence warnings directly via idata.sample_stats.warning, at the cost of not being able to save these InferenceData objects.

Docs / Maintenance

• New pm.stats.convergence submodule for functional implementation of convergence diagnostics and warnings.

[codecov]

Codecov Report

Merging #6192 (4096bea) into main (570e6e8) will decrease coverage by 0.03%.
The diff coverage is 84.41%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6192      +/-   ##
==========================================
- Coverage   93.77%   93.74%   -0.04%     
==========================================
  Files         101      101              
  Lines       22232    22209      -23     
==========================================
- Hits        20849    20819      -30     
- Misses       1383     1390       +7     

Impacted Files	Coverage Δ
pymc/backends/base.py	88.47% <ø> (+0.52%)	⬆️
pymc/step_methods/compound.py	84.61% <ø> (-2.06%)	⬇️
pymc/parallel_sampling.py	88.05% <60.00%> (+2.52%)	⬆️
pymc/step_methods/hmc/base_hmc.py	89.81% <75.00%> (-0.74%)	⬇️
pymc/stats/convergence.py	90.90% <85.71%> (-2.25%)	⬇️
pymc/sampling.py	83.58% <94.44%> (+1.04%)	⬆️
pymc/step_methods/hmc/hmc.py	92.72% <100.00%> (+0.13%)	⬆️
pymc/step_methods/hmc/nuts.py	97.26% <100.00%> (-0.15%)	⬇️
pymc/tests/step_methods/hmc/test_nuts.py	100.00% <100.00%> (ø)
pymc/step_methods/step_sizes.py	75.00% <0.00%> (-25.00%)	⬇️
... and 4 more

[aseyboldt]

Hm, I really don't like loosing the extra divergence info. I'm not sure how many people actually use this, but I certainly do :-)
If we want to get rid of the warning system, I think that's fine, but we should then store the information in in the sampler stats instead (or at least have an option to do that).

[michaelosthege]

Hm, I really don't like loosing the extra divergence info. I'm not sure how many people actually use this, but I certainly do :-) If we want to get rid of the warning system, I think that's fine, but we should then store the information in in the sampler stats instead (or at least have an option to do that).

Unfortunately it's not that easy, because a stat can't be a dictionary.
In addition it is not clear whether stats can be sparse.

Assuming we want to store them for later it would be good to use the same data structure across HMC implementations (PyMC, nutpie, numpyro). Do you have any idea which structure that could be?

[michaelosthege]

I think we can compromise:
Storing SamplerWarning-typed items in the stats dict and idata.sample_stats group actually works.
...until one tries to save it to_netcdf.

Coordinating with Oriol, I'll open a PR in ArviZ to filter out object-typed variables from the groups before saving. This will make to_netcdf also more robust in case there are other variables with invalid dtype.
(We can try converting them, but if this fails log a warning..)

[aseyboldt]

I think we should be able to come up with a nice format for those in the sampler stats, without resorting to storing objects.
The most tricky part is dealing with the sparse structure, but we could handle that with a hierarchical index, or alternatively a second variable that maps the divergence to the chain. So something along the lines of this:

# Extra dims in the sample_stats:
"divergence" # One key for each divergence across all chains
"unconstrained_parameter"
# Extra vars:
"divergence_start": ("divergence", "unconstrained_parameter")
"divergence_end": ("divergence", "unconstrained_parameter")
"divergence_chain": ("divergence")  # For each divergence the label of the chain where that divergence happend

Instead of "divergence_chain" we could also use a hierarchical index for the divergence index. I think there is still some trouble when storing hierarchical indices as netcdf, but maybe after an unstack that would be fine?

[michaelosthege]

@aseyboldt a serializable structure would be nice, yes, but I see other issues with higher priority..
With arviz-devs/arviz#2134 ArviZ will automatically drop object variables before saving which means that we can safely have the SamplerWarning end up InferenceData. It just won't be saved, but that wasn't the case before either.

To clean up the sampler stats structure, IMO the next step is #6207 and introducing a coord/naming scheme by which samplers can be identified.

This should also open the door to storing other sampler-wise information such as which variables they were sampling.

In the meantime, this PR will have to wait for the upcoming ArviZ release.

[michaelosthege]

The last changes I pushed made thus PR independent of ArviZ updates.

I updated the description, so this is ready to review/merge :)

[michaelosthege]

@pymc-devs/core-contributors can I get reviews here?
I'd like to move forward before git conflicts creep in..

[OriolAbril]

I think we should be able to come up with a nice format for those in the sampler stats, without resorting to storing objects.
The most tricky part is dealing with the sparse structure, but we could handle that with a hierarchical index, or alternatively a second variable that maps the divergence to the chain. So something along the lines of this:

# Extra dims in the sample_stats:
"divergence" # One key for each divergence across all chains
"unconstrained_parameter"
# Extra vars:
"divergence_start": ("divergence", "unconstrained_parameter")
"divergence_end": ("divergence", "unconstrained_parameter")
"divergence_chain": ("divergence")  # For each divergence the label of the chain where that divergence happend

I think it is worth it to move that to an issue. I tried going over the data in the divergences warning but was unable to understand what info was exactly there and how to store it in "regular" arrays. Defining a structure that works in all cases might require a bit of extra work to account for multiple samplers or adding some exceptions for this to work if nuts is the only sampler, but it might be somewhat feasible. At least on xarray side, the dimensions are shared, but variables in a dataset can have any dimensions (there is not need for all variables to have chain and draw dims for example) and a dataset can have as many dimensions and coordinates as desired.

[lucianopaz]

Looks good to me, @michaelosthege . I agree with @OriolAbril that we need to open an issue to change the warning data structure to something that is serializable. It’s a real waste to lose all the warning information when storing the results to netcdf.

[michaelosthege]

Thanks Luciano!

As you can see, I created #6252.

I'll merge this before another git conflict creeps in..