Vertical grouping of labels with brackets on matplotlib

Question

I want to display vertical brackets to group ylabels on the left side of a plot, to group for instance the first 3 variables together, then the following two.

By carefully tuning the position of xy and xytest of annotate using fancyarrow I can manage to get somehow vertical brackets using annotate and fancyarrows but placing them without breaking the alignment is such a hassle that it's near undoable. This is because the process is at the same time extremely painful and brittle because one cannot just place xy and xytest at the same horizontal, because the angle depends on the text size and the length of the bracket. Plus with the rotation it's hard to manage the direction of the bracket on my example it's reversed for some reasons.

Here is my very poor unsuccessful trial:

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from matplotlib.pyplot import figure
figure(figsize=(10, 5), dpi=50)

np.random.seed(42)
df = pd.DataFrame({f"var{i}": np.random.uniform(-1., 1., 2) for i in range(10)})
df["facet"] = [0, 1]
# Vertical plot
df = df.melt('facet', var_name='Variable',  value_name='values')
ax = sns.barplot(data=df, x="values", y="Variable")

ylims = ax.get_ylim()
xlims = ax.get_xlim()
xmin = min(xlims)
xmax = max(xlims)
ymin = min(ylims)
ymax = max(ylims)

ax.annotate('Group1 vertical', xytext=(1.4 * xmin, ymax / 2), xy=(1.2 * xmin, ymax / 2 - (ymax - ymin)/7.5),
            fontsize=10, ha='center', va='bottom',
            bbox=dict(boxstyle='square', fc='white'),
            arrowprops=dict(arrowstyle=f'-[, widthB=10, lengthB={0.5 * xmin}', lw=2.0), rotation=90, clip_on=False, annotation_clip=False)

ax.annotate('Group2 vertical', xytext=(1.4 * xmin, ymax * 4 / 5), xy=(1.2 * xmin, ymax * 4 / 5 - (ymax - ymin)/7.5),
            fontsize=10, ha='center', va='bottom',
            bbox=dict(boxstyle='square', fc='white'),
            arrowprops=dict(arrowstyle=f'-[, widthB=5, lengthB={0.5 * xmin}', lw=2.0), rotation=90, clip_on=False, annotation_clip=False)

ax.yaxis.label.set_visible(False)
ax.xaxis.label.set_visible(False)
plt.tight_layout()
plt.show()

Can it be done and done somehow elegantly once you have placed xy or xytest for instance and fixed the length and width of the bracket ?

user_380231 · Answer

One way to do this is to make use of the get_yaxis_transform method which returns as blended transform where the x-values are in "Axes Fraction" and the y-value are in "data". Once we have that transform we can create a PathPatch (docs) for the bracket and use ax.text (docs) to add the label. I wrapped that logic up in a function to make it easier to use it a couple of times:

import matplotlib.path as mpath
import matplotlib.patches as mpatches


def add_label_band(ax, top, bottom, label, *, spine_pos=-0.05, tip_pos=-0.02):
    """
    Helper function to add bracket around y-tick labels.

    Parameters
    ----------
    ax : matplotlib.Axes
        The axes to add the bracket to

    top, bottom : floats
        The positions in *data* space to bracket on the y-axis

    label : str
        The label to add to the bracket

    spine_pos, tip_pos : float, optional
        The position in *axes fraction* of the spine and tips of the bracket.
        These will typically be negative

    Returns
    -------
    bracket : matplotlib.patches.PathPatch
        The "bracket" Aritst.  Modify this Artist to change the color etc of
        the bracket from the defaults.

    txt : matplotlib.text.Text
        The label Artist.  Modify this to change the color etc of the label
        from the defaults.

    """
    # grab the yaxis blended transform
    transform = ax.get_yaxis_transform()

    # add the bracket
    bracket = mpatches.PathPatch(
        mpath.Path(
            [
                [tip_pos, top],
                [spine_pos, top],
                [spine_pos, bottom],
                [tip_pos, bottom],
            ]
        ),
        transform=transform,
        clip_on=False,
        facecolor="none",
        edgecolor="k",
        linewidth=2,
    )
    ax.add_artist(bracket)

    # add the label
    txt = ax.text(
        spine_pos,
        (top + bottom) / 2,
        label,
        ha="right",
        va="center",
        rotation="vertical",
        clip_on=False,
        transform=transform,
    )

    return bracket, txt

And to use it:


import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from matplotlib.pyplot import figure

figure(figsize=(10, 5), dpi=50)

np.random.seed(42)
df = pd.DataFrame({f"var{i}": np.random.uniform(-1.0, 1.0, 2) for i in range(10)})
df["facet"] = [0, 1]
# Vertical plot
df = df.melt("facet", var_name="Variable", value_name="values")
ax = sns.barplot(data=df, x="values", y="Variable")

ylims = ax.get_ylim()
xlims = ax.get_xlim()
xmin = min(xlims)
xmax = max(xlims)
ymin = min(ylims)
ymax = max(ylims)

ax.yaxis.label.set_visible(False)
ax.xaxis.label.set_visible(False)

add_label_band(ax, 0.75, 5.25, "group a")
add_label_band(ax, 6.75, 8.25, "group b")

plt.tight_layout()
plt.show()

This could probably use a bit more tweaking to get the values to be more visually pleasing and you could extend the signature to pass styling through to the bracket and text.

You could also replace the ax.text call with ax.annotate if you want the line, but I was not sure if that was a key feature of the OP or not.

Vertical grouping of labels with brackets on matplotlib

1 Answers