assign reference values by var_name instead of var_label

Author: tvwenger

arviz/plots/backends/bokeh/pairplot.py

@@ -37,6 +37,7 @@ def plot_pair(
     diverging_mask,
     divergences_kwargs,
     flat_var_names,
+    flat_var_labels,
     backend_kwargs,
     marginal_kwargs,
     show,
@@ -262,8 +263,8 @@ def get_width_and_height(jointplot, rotate):
                     **marginal_kwargs,
                 )
 
-                ax[j, i].xaxis.axis_label = flat_var_names[i]
-                ax[j, i].yaxis.axis_label = flat_var_names[j + marginals_offset]
+                ax[j, i].xaxis.axis_label = flat_var_labels[i]
+                ax[j, i].yaxis.axis_label = flat_var_labels[j + marginals_offset]
 
             elif j + marginals_offset > i:
                 if "scatter" in kind:
@@ -350,8 +351,8 @@ def get_width_and_height(jointplot, rotate):
                     y = reference_values_copy[flat_var_names[i]]
                     if x and y:
                         ax[j, i].scatter(y, x, **reference_values_kwargs)
-                ax[j, i].xaxis.axis_label = flat_var_names[i]
-                ax[j, i].yaxis.axis_label = flat_var_names[j + marginals_offset]
+                ax[j, i].xaxis.axis_label = flat_var_labels[i]
+                ax[j, i].yaxis.axis_label = flat_var_labels[j + marginals_offset]
 
     show_layout(ax, show)
 

arviz/plots/backends/matplotlib/pairplot.py

@@ -30,6 +30,7 @@ def plot_pair(
     diverging_mask,
     divergences_kwargs,
     flat_var_names,
+    flat_var_labels,
     backend_kwargs,
     marginal_kwargs,
     show,
@@ -215,8 +216,8 @@ def plot_pair(
                 reference_values_copy[flat_var_names[1]],
                 **reference_values_kwargs,
             )
-        ax.set_xlabel(f"{flat_var_names[0]}", fontsize=ax_labelsize, wrap=True)
-        ax.set_ylabel(f"{flat_var_names[1]}", fontsize=ax_labelsize, wrap=True)
+        ax.set_xlabel(f"{flat_var_labels[0]}", fontsize=ax_labelsize, wrap=True)
+        ax.set_ylabel(f"{flat_var_labels[1]}", fontsize=ax_labelsize, wrap=True)
         ax.tick_params(labelsize=xt_labelsize)
 
     else:
@@ -344,12 +345,12 @@ def plot_pair(
                 if j != vars_to_plot - 1:
                     plt.setp(ax[j, i].get_xticklabels(), visible=False)
                 else:
-                    ax[j, i].set_xlabel(f"{flat_var_names[i]}", fontsize=ax_labelsize, wrap=True)
+                    ax[j, i].set_xlabel(f"{flat_var_labels[i]}", fontsize=ax_labelsize, wrap=True)
                 if i != 0:
                     plt.setp(ax[j, i].get_yticklabels(), visible=False)
                 else:
                     ax[j, i].set_ylabel(
-                        f"{flat_var_names[j + not_marginals]}",
+                        f"{flat_var_labels[j + not_marginals]}",
                         fontsize=ax_labelsize,
                         wrap=True,
                     )

arviz/plots/pairplot.py

@@ -196,7 +196,8 @@ def plot_pair(
             get_coords(dataset, coords), var_names=var_names, skip_dims=combine_dims, combined=True
         )
     )
-    flat_var_names = [
+    flat_var_names = [var_name for var_name, _, _, _ in plotters]
+    flat_var_labels = [
         labeller.make_label_vert(var_name, sel, isel) for var_name, sel, isel, _ in plotters
     ]
 
@@ -253,6 +254,7 @@ def plot_pair(
         diverging_mask=diverging_mask,
         divergences_kwargs=divergences_kwargs,
         flat_var_names=flat_var_names,
+        flat_var_labels=flat_var_labels,
         backend_kwargs=backend_kwargs,
         marginal_kwargs=marginal_kwargs,
         show=show,