hierarchical plotting functionality

[chrberrig]

Hi all.

I note from working with the different plotting functionality, that hierarchies are not entirely without problems in arviz.

I wonder if this can somehow be implemented?

For what its worth, i have done some thinking regarding this subject:

First, the information is available though the coordinates of any particular dimension/coords declaration (for example, a hierarchy over geography, can be done by finest resolution, say region, but to each region coordinate, an associated assigned coordinate could assign the country, one level higher in the hierarchy: this is simply put a wide-format coordinate set depending on the level in the hierarchy, not an index and not an dimension since its simply a nested unit. ).

So implementation-wise the solution does not seem far away, but most plotting functions need handles for this.

I guess that what I'm asking about, is how this would be handled from the "grammar of graphics" angle, as this seem to be the handling the figures aim for in general.

As a continuation of the example above, plot_ridge could relatively easily do this in theory:
Two levels of parameters exists b_region, and b_country on level above.
this could be "solved" by allowing hierarchical ordering for labels across variables instead of having variables as foundation for anything else in this plot, as well as allowing the access to the second level of the region coordinates.

What I'm not saying is that there exists a universal handle for hierarchical models, but rather that hierarchical handling should exist, since the information is easily encoded though additional coordinate assignment when loading fits into datatrees anyway.

[chrberrig]

Concrete case: aes mapped to a non-dimension coordinate breaks target resolution

Following up on the general point above with a specific, reproducible failure.

Setup: A hierarchical model where roder (44 neighbourhoods) nest within distrikter (12 districts). The parameter b_age_rode has dims (age_group, rode) — there is no distrikt dimension, because the nesting is a relation between index sets, not an array axis. Following the Stan model's distrikt_of[AR] array, I attach it as a coordinate:

ds = ds.assign_coords(distrikt_of_rode=("rode", rode_to_district))

I then want the 44 rode ridges coloured by their district.

What works: generate_aes_dt accepts the coordinate — it checks membership with set(da.dims).union(da.coords), and the size logic correctly collapses 44 values to 12 uniques:

unique_values = np.unique(data[dim])
clean_sizes[dim] = len(unique_values)

What fails: In map:

dim_to_idx = {data[idx].dims[0]: idx for idx in data.coords
              if (idx in all_loop_dims) and (idx not in data.dims)}
# -> {"rode": "distrikt_of_rode"}
plotters = xarray_sel_iter(data, skip_dims=skip_dims, dim_to_idx=dim_to_idx)

dim_to_idx substitutes the dimension, so iteration now yields one selection per district (12) rather than per rode (44). Faceting is still per rode, so get_target returns a DataArray of axes rather than one, and backend_from_object raises:

ValueError: Attempting to get plotting backend from a non-scalar object.
Conflicting xarray object:
<xarray.DataArray 'b_age_rode' (rode: 4)> array([<Axes>, <Axes>, <Axes>, <Axes>])
Coordinates:
  * rode              (rode) object 'Amalienborgs-Toldbodens Rode' ...
    distrikt_of_rode  (rode) object 'Sankt Annæ Østre Kvarter' ...

Four axes, all from one district — the loop is at district granularity while the target lookup is at rode granularity.

Reproduction: (arviz-plots 1.2.0, matplotlib backend):

#import numpy as np, xarray as xr, arviz_plots as azp

rng = np.random.default_rng(0)
n_group, n_item = 3, 9
parent = np.repeat([f"g{i}" for i in range(n_group)], n_item // n_group)

ds = xr.Dataset(
    {"theta": (("chain", "draw", "item"), rng.normal(size=(4, 200, n_item)))},
    coords={"item": [f"i{j}" for j in range(n_item)]},
).assign_coords(parent=("item", parent))

azp.plot_dist(
    ds, 
    aes={"color": ["parent"]}, 
)   # raises

Same failure in plot_ridge and plot_forest.

Question: Should dim_to_idx retain the underlying dimension in sel alongside the index coordinate, so aesthetics are looked up by coordinate while iteration and target resolution stay at the dimension's granularity? Colouring 44 items by 12 groups seems like it should keep 44 iterations, not collapse to 12.

If that's the intended direction I'm happy to attempt a PR — pointers on whether the fix belongs in map, in get_target, or in xarray_sel_iter would help.

[chrberrig]

Filed the concrete case as arviz-plots issue #551 — coordinate-based aesthetics fail on target resolution.