Skip to content

Fix dimension labelling of plot_ridge: - #507

Open
chrberrig wants to merge 1 commit into
arviz-devs:mainfrom
chrberrig:fix_dim_labeller_ridge_plot
Open

Fix dimension labelling of plot_ridge: #507
chrberrig wants to merge 1 commit into
arviz-devs:mainfrom
chrberrig:fix_dim_labeller_ridge_plot

Conversation

@chrberrig

Copy link
Copy Markdown
Contributor

This is done by making dim_map for MapLabeller passable in the labeller kwarg for plot_ridge.

Added tests for this, which ran successfully

Also did local rendering in notebook, which works.

… for MapLabeller passable in the labeller kwarg for plot_ridge. Also added tests for this
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87.95%. Comparing base (f386e45) to head (a4477d8).

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #507      +/-   ##
==========================================
+ Coverage   87.87%   87.95%   +0.08%     
==========================================
  Files          63       63              
  Lines        7025     7025              
==========================================
+ Hits         6173     6179       +6     
+ Misses        852      846       -6     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@read-the-docs-community

Copy link
Copy Markdown

[label.strip("_") for label in labellable_dims if label in labels],
[
(
labeller.var_name_to_str(label.strip("_"))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be calling to dim_coord_to_str(label.strip..., {}, {}) instead which will automatically handle the dim map case if the input is a MapLabeller but would also work for any other user defined labeller that processes the dimension names.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried replacing direct dim_map access with labeller.dim_coord_to_str, but BaseLabeller.dim_coord_to_str returns the coordinate value, not the dimension string, so calling it for a dimension header gives None/empty output unless a coordinate value is supplied.

For the tick label header we only have the dimension name, not a coordinate value. Is there an existing public labeller method intended to map dimension names alone, or should I add a small helper/fallback here?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I could imagine an approach where we define a BaseLabeller.dim_name_to_str method that mirrors var_name_to_str, but for dimension labelling, and then pass that through a mapped dimension name in MapLabeller.

I realize this may overlap somewhat with the intended role of dim_coord_to_str, so this might not be the right direction. But from what I can tell, BaseLabeller.dim_coord_to_str currently only returns coordinate label strings, which means the dimension information gets lost when calling super() from MapLabeller, even though the dimension mapping is still available locally in MapLabeller.

@OriolAbril: you probably know the codebase much better than I do, but I tried tracing where BaseLabeller.dim_coord_to_str is used, and I couldn’t find any usage outside of the MapLabeller super() call to BaseLabeller. Would it make sense to adjust the return value of dim_coord_to_str at the BaseLabeller level somehow so the dimension information can be preserved OR to have MapLabeller overwrite the method from super()? I’m not entirely sure what would fit best with the existing design patterns in the library, nor how this adjustment would look exactly, but I’d be curious to hear your thoughts.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be calling to dim_coord_to_str(label.strip..., {}, {}) instead which will automatically handle the dim map case if the input is a MapLabeller but would also work for any other user defined labeller that processes the dimension names.

I looked into dim_coord_to_str and mix_labellers. My understanding is that they operate on dimension-coordinate pairs, not on dimension names alone. In BaseLabeller, dim_coord_to_str returns the coordinate value, while MapLabeller maps dim/coord_val and delegates to super(). That works well for selections, but the tick label here is a dimension-name header without a coordinate value.

So I don’t currently see a public labeller method that maps a dimension name alone. The options I see are either keeping a guarded getattr(labeller, "dim_map", {}) fallback for MapLabeller, or adding a public dim_name_to_str method in arviz-base. Am I missing an existing method that is intended for this?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I opened an arviz-base PR adding a public dim_to_str method that can be used here instead of accessing dim_map directly. See Add dim_to_str method to labellers #200 of arviz base.
This can be used for a robust labeller method for this fix.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I understand the intended direction now: use dim_coord_to_str and make the plot use a dimension-aware labeller such as DimCoordLabeller when dimension names are needed.

One thing I want to confirm: with a plain MapLabeller(dim_map={"hierarchy": "MyHierarchy"}), dim_coord_to_str("hierarchy", "", "") still delegates to BaseLabeller, which omits the dimension name, so it does not produce "MyHierarchy". To get "MyHierarchy", the user would need a mixed labeller like mix_labellers((MapLabeller, DimCoordLabeller)).

Is that the intended user-facing behavior for this case?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Im sorry; many of my previous comments got stuck as reviews, and i did not realize you could not see them: please let me know what you think, and sorry for the wall of text !

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I'm starting to understand the intended direction better now.

I experimented with using dim_coord_to_str for the header labels and was able to make the dimension header pick up the mapping when using a mixed labeller such as:

mix_labellers((MapLabeller, DimCoordLabeller))

However, while doing so I noticed something that makes me hesitate.

The plot seems to have two distinct labeling contexts:

  1. Row/value labels, which are already handled through the existing labeller machinery (make_label_flat, sel_to_str, etc.).
  2. Column/header labels, which are derived from labels / labellable_dims.

For the second case, the labels are dimension or pseudo-dimension names (__variable__, school, hierarchy, ...), not dimension-coordinate pairs.

Using dim_coord_to_str for the headers does make the header mapping work, but only through a dimension-aware labeller. More importantly, the same mixed labeller also changes the coordinate labels throughout the plot. For example, instead of:

Choate
Deerfield
...

the labels become:

MySchool: Choate
MySchool: Deerfield
...

which is the correct behaviour of DimCoordLabeller, but is different from the original plot output.

So the implementation works, but it also seems to highlight that the concepts "dimension header" and "dimension-coordinate label" are not necessarily the same thing.

Am I understanding the intended behaviour correctly? Should the coordinate labels also change in this situation, or is there another direction you had in mind?

Just for reference, the changed code regarding this behaviour is

if ticklabel_kwargs is not False:    
┃         def _label_header_to_str(label, labeller):    
┃             if label == "__variable__":    
┃                 return labeller.var_name_to_str(label.strip("_"))    
┃     
┃             dim_label = labeller.dim_coord_to_str(label, "", "").rstrip(": ")    
┃             if dim_label:    
┃                 return dim_label    
┃               Not Committed Yet    
┃             return label     
┃     
┃         labels_list = [    
┃             _label_header_to_str(label, labeller)     
┃             for label in labellable_dims if label in labels    
┃         ]    
┃         print(labels_list)    
          plot_bknd.xticks(    
              np.arange(len(labels)),    
~             labels_list,     
              plot_collection.get_target(None, {"column": "labels"}),    
              **ticklabel_kwargs,    
          )    

and the generated graphic is
Screenshot From 2026-06-14 10-38-49

with the notebook settings of

azp.style.use("arviz-variat")

non_centered = azb.load_arviz_data('non_centered_eight')

var_name_map={
    "theta": r"$\theta$", 
    "theta_t": r"$\theta_{t}$", 
    "mu": r"$\mu$", 
    "tau": r"$\tau$", 
}

dim_map={
    # "school": r"$\Sigma$", 
    "school": "MySchool", 
}

labeller = azb.labels.MapLabeller(
    var_name_map=var_name_map, 
    dim_map=dim_map, 
)

LabellerCustom = azb.labels.mix_labellers(
    (azb.labels.MapLabeller, azb.labels.DimCoordLabeller), 
)
labeller_local = LabellerCustom(
    var_name_map=var_name_map, 
    dim_map=dim_map, 
)

pc = azp.plot_ridge(
    non_centered,
    var_names=["theta", "mu", "theta_t", "tau"],
    aes={"color": ["__variable__"]},
    figure_kwargs={"width_ratios": [1, 2], "layout": "none"},
    aes_by_visuals={"labels": ["color"]},
    shade_label="school",
    labeller=labeller_local, 
    # labels=["__variable__"], #, "school"],
    # labels=["school"],
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I completely missed that annotate_label, the function in the visuals module used for the coordinate annotations, uses make_label_flat and relies on getting no var_name to skip that portion of the label instead of directly calling the lower level method. That makes the DimCoordLabeller not work as the default as you have seen. We might need to both add some extra methods to the labeller classes and think about how the labeller is used in plot_forest/ridge so we actually use all the granularity available. Let me think about it and get back, thanks for the patience and properly looking into the behaviour.

@chrberrig chrberrig Jun 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No worries! thanks for getting back!

I were actually thinking about whither a wider discussion on the labeller API would be prudent?

It would appear that the current labeller API appears to have explicit concepts for:

  • variable names
  • coordinate values
  • dimension-coordinate labels

while the labels used as plot headers seem to be a separate concept:

  • variable-dimension labels
  • dimension headers/labels
  • pseudo-dimension headers/labels (such as variable)

I simply have a suspicion that there is a gap in the current API in this regard, but I'm not sure if (or where) such a discussion would be appropriate to have?

let me know what you think, and thanks for getting back to me so fast :)

Comment thread tests/test_plots.py
Comment on lines +963 to +972
def _get_ticklabel_texts(plot, backend):
if backend == "matplotlib":
return [tick.get_text() for tick in plot.get_xticklabels()]
if backend == "bokeh":
return list(plot.xaxis[0].major_label_overrides.values())
if backend == "plotly":
return list(plot.layout.xaxis.ticktext)
if backend == "none":
return next(item["labels"] for item in plot if item.get("function") == "xticks")
raise ValueError(f"Unknown backend: {backend}")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is a single function, not one per backend, so I don't see the need for this. I think checking the labels on the none backend is more than enough. We could add these kind of tests on a TestPlotLabels class that is not parametrized by backend.

@OriolAbril OriolAbril mentioned this pull request Jun 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants