Annotating ranges of data in matplotlib

As mentioned in this answer, you can construct curly brackets with sigmoidal functions. Below is a function that adds curly brackets just above the x-axis. The curly brackets it produces should look the same regardless of the axes limits, as long as the figure width and height don't vary.

import numpy as np
import matplotlib.pyplot as plt

def draw_brace(ax, xspan, text):
    """Draws an annotated brace on the axes."""
    xmin, xmax = xspan
    xspan = xmax - xmin
    ax_xmin, ax_xmax = ax.get_xlim()
    xax_span = ax_xmax - ax_xmin
    ymin, ymax = ax.get_ylim()
    yspan = ymax - ymin
    resolution = int(xspan/xax_span*100)*2+1 # guaranteed uneven
    beta = 300./xax_span # the higher this is, the smaller the radius

    x = np.linspace(xmin, xmax, resolution)
    x_half = x[:resolution//2+1]
    y_half_brace = (1/(1.+np.exp(-beta*(x_half-x_half[0])))
                    + 1/(1.+np.exp(-beta*(x_half-x_half[-1]))))
    y = np.concatenate((y_half_brace, y_half_brace[-2::-1]))
    y = ymin + (.05*y - .01)*yspan # adjust vertical position

    ax.autoscale(False)
    ax.plot(x, y, color='black', lw=1)

    ax.text((xmax+xmin)/2., ymin+.07*yspan, text, ha='center', va='bottom')

ax = plt.gca()
ax.plot(range(10))
draw_brace(ax, (0, 8), 'large brace')
draw_brace(ax, (8, 9), 'small brace')

Output:

enter image description here


I modified Joooeey's answer to allow to change the vertical position of braces:

def draw_brace(ax, xspan, yy, text):
    """Draws an annotated brace on the axes."""
    xmin, xmax = xspan
    xspan = xmax - xmin
    ax_xmin, ax_xmax = ax.get_xlim()
    xax_span = ax_xmax - ax_xmin

    ymin, ymax = ax.get_ylim()
    yspan = ymax - ymin
    resolution = int(xspan/xax_span*100)*2+1 # guaranteed uneven
    beta = 300./xax_span # the higher this is, the smaller the radius

    x = np.linspace(xmin, xmax, resolution)
    x_half = x[:int(resolution/2)+1]
    y_half_brace = (1/(1.+np.exp(-beta*(x_half-x_half[0])))
                    + 1/(1.+np.exp(-beta*(x_half-x_half[-1]))))
    y = np.concatenate((y_half_brace, y_half_brace[-2::-1]))
    y = yy + (.05*y - .01)*yspan # adjust vertical position

    ax.autoscale(False)
    ax.plot(x, y, color='black', lw=1)

    ax.text((xmax+xmin)/2., yy+.07*yspan, text, ha='center', va='bottom')
ax = plt.gca()
ax.plot(range(10))
draw_brace(ax, (0, 8), -0.5, 'large brace')
draw_brace(ax, (8, 9), 3, 'small brace')

Output:

matplotlib braces

Also note that in Joooeey's answer, line

x_half = x[:resolution/2+1]

should be

x_half = x[:int(resolution/2)+1]

Otherwise, the number that the script tries to use as index here is a float.

Finally, note that right now the brace will not show up if you move it out of bounds. You need to add parameter clip_on=False, like this:

ax.plot(x, y, color='black', lw=1, clip_on=False)

You can just wrap it all up in a function:

def add_range_annotation(ax, start, end, txt_str, y_height=.5, txt_kwargs=None, arrow_kwargs=None):
    """
    Adds horizontal arrow annotation with text in the middle

    Parameters
    ----------
    ax : matplotlib.Axes
        The axes to draw to

    start : float
        start of line

    end : float
        end of line

    txt_str : string
        The text to add

    y_height : float
        The height of the line

    txt_kwargs : dict or None
        Extra kwargs to pass to the text

    arrow_kwargs : dict or None
        Extra kwargs to pass to the annotate

    Returns
    -------
    tuple
        (annotation, text)
    """

    if txt_kwargs is None:
        txt_kwargs = {}
    if arrow_kwargs is None:
        # default to your arrowprops
        arrow_kwargs = {'arrowprops':dict(arrowstyle="<->",
                            connectionstyle="bar",
                            ec="k",
                            shrinkA=5, shrinkB=5,
                            )}

    trans = ax.get_xaxis_transform()

    ann = ax.annotate('', xy=(start, y_height),
                        xytext=(end, y_height),
                        transform=trans,
                        **arrow_kwargs)
    txt = ax.text((start + end) / 2,
                  y_height + .05,
                  txt_str,
                  **txt_kwargs)


    if plt.isinteractive():
        plt.draw()
    return ann, txt

Alternately,

start, end = .6, .8
ax.axvspan(start, end, alpha=.2, color='r')
trans = ax.get_xaxis_transform()
ax.text((start + end) / 2, .5, 'test', transform=trans)

Tags:

Matplotlib