Strange behaviour of the legend of plot_trace when using coords

[lciti]

Describe the bug
When using plot_trace with idata generated using a model with coordinates, the legend does not reflect the coordinate names but something completely different (the chains line types, even if chains are not plotted with combined=True), at least for the first variable plotted.

To Reproduce

with pm.Model(coords={'a': [1, 2], 'b': ['+', '-']}) as m:
    pm.Normal('RV', dims=['a', 'b'])
    idata = pm.sample()
az.plot_trace(idata, combined=True, legend=True)

However, it seems to work fine for the second or later variables plotted.

with pm.Model(coords={'a': [1, 2], 'b': ['+', '-']}) as m:
    pm.Normal('RV0', dims=['a', 'b'])
    pm.Normal('RV', dims=['a', 'b'])
    idata = pm.sample()
az.plot_trace(idata, combined=True, legend=True)

Expected behavior
The expected behaviour when combined=True is the one shown for the second variable in the second example, where the legend shows "1, +", "1, -", etc.

Additional context
Watermark:

Last updated: Mon Mar 11 2024

Python implementation: CPython
Python version       : 3.8.2
IPython version      : 8.2.0

xarray: 0.20.1

graphviz: 0.16
arviz   : 0.13.0
numpy   : 1.24.4
pandas  : 1.4.2
pytensor: 2.12.3
pymc    : 5.6.1

Watermark: 2.4.3

[imperorrp]

I tried recreating this with more recent versions-

Python 3.11.8
arviz==0.17.1
matplotlib==3.7.1
pymc==5.11.0

And got exactly the same plot outputs as you describe here. I can try and see if I can find out why this might be happening.

[OriolAbril]

The legend you see in the 2nd plot is always generated, but there is this bug that it is sometimes overwriten by the chain one. If the chain legend was moved to the right column with the trace itself instead of the densities I think this would be fixed.

[imperorrp]

The legend you see in the 2nd plot is always generated, but there is this bug that it is sometimes overwriten by the chain one. If the chain legend was moved to the right column with the trace itself instead of the densities I think this would be fixed.

This only happens sometimes? I could try making that change you suggest though(moving the chain legend to the right column with the trace)- will create a PR if it works!

[imperorrp]

Just linked a PR for this!

Very small edit, but hope it helps