Add include_transformed option when computing ESS

[ricardoV94]

I am moving this old pymc3 issue pymc-devs/pymc#3626 here.

That issue asked for a flag to include/exclude transformed_variables when computing the ESS. My understanding is that Arviz does not really distinguish between those natively or does it? Also I am not sure how the ESS is computed, but wouldn't it be the same for transformed and the untransformed version?

Similarly this already closed pymc3 issue raised the same issue: pymc-devs/pymc#3420 in relation to the traceplot.

Just to avoid confusion, this does not correspond to the transform argument that you can pass to most plots to transform the data before displaying it. This is about including or excluding automatically transformed variables from pymc such as a_log__ ~ TransformedDistribution that are created when calling a = pm.HalfNormal('a').

[AlexAndorra]

Thanks for transferring the issue @ricardoV94 !

Yeah, I don't know how this relates to ESS either, but I can see that being quite useful when exploring a model with divergences, as it's usually better to examine the parameters in the space that the sampler uses.

We talked about that on the ArviZ Slack (cc @ahartikainen @aloctavodia and @ColCarroll), and a potentially nice way would be to store the transformations into a dictionary mapping each RV to its transformation, e.g {tau: log, alpha: logodds} in, for instance, a new InferenceData group. And then you could probably pass that to the transform kwarg that most plots have!

[aloctavodia]

The ESS and Rhat diagnostics should be the same for transformed and untransformed variables.

[aloctavodia]

BTW this meta-data could also help to identify circular variables