Fix dimension labelling of plot_ridge: - #507
Conversation
… for MapLabeller passable in the labeller kwarg for plot_ridge. Also added tests for this
Codecov Report✅ All modified and coverable lines are covered by tests. 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. 🚀 New features to boost your workflow:
|
Documentation build overview
60 files changed ·
|
| [label.strip("_") for label in labellable_dims if label in labels], | ||
| [ | ||
| ( | ||
| labeller.var_name_to_str(label.strip("_")) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
This should be calling to
dim_coord_to_str(label.strip..., {}, {})instead which will automatically handle the dim map case if the input is aMapLabellerbut 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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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 !
There was a problem hiding this comment.
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:
- Row/value labels, which are already handled through the existing labeller machinery (
make_label_flat,sel_to_str, etc.). - 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,
)
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"],
)
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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 :)
| 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}") |
There was a problem hiding this comment.
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.

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.