Skip to content

Instantly share code, notes, and snippets.

@lzkelley
Last active July 16, 2020 15:35
Show Gist options
  • Save lzkelley/0de9e8bf2a4fe96d2018f1b1bd5a0d3c to your computer and use it in GitHub Desktop.
Save lzkelley/0de9e8bf2a4fe96d2018f1b1bd5a0d3c to your computer and use it in GitHub Desktop.
Add an annotation (label) to a matplotlib line. Position and Alignment is adjusted to match the line.
def label_line(ax, line, label, color='0.5', fs=14, halign='left'):
"""Add an annotation to the given line with appropriate placement and rotation.
Based on code from:
[How to rotate matplotlib annotation to match a line?]
(http://stackoverflow.com/a/18800233/230468)
User: [Adam](http://stackoverflow.com/users/321772/adam)
Arguments
---------
ax : `matplotlib.axes.Axes` object
Axes on which the label should be added.
line : `matplotlib.lines.Line2D` object
Line which is being labeled.
label : str
Text which should be drawn as the label.
...
Returns
-------
text : `matplotlib.text.Text` object
"""
xdata, ydata = line.get_data()
x1 = xdata[0]
x2 = xdata[-1]
y1 = ydata[0]
y2 = ydata[-1]
if halign.startswith('l'):
xx = x1
halign = 'left'
elif halign.startswith('r'):
xx = x2
halign = 'right'
elif halign.startswith('c'):
xx = 0.5*(x1 + x2)
halign = 'center'
else:
raise ValueError("Unrecogznied `halign` = '{}'.".format(halign))
yy = np.interp(xx, xdata, ydata)
ylim = ax.get_ylim()
# xytext = (10, 10)
xytext = (0, 0)
text = ax.annotate(label, xy=(xx, yy), xytext=xytext, textcoords='offset points',
size=fs, color=color, zorder=1,
horizontalalignment=halign, verticalalignment='center_baseline')
sp1 = ax.transData.transform_point((x1, y1))
sp2 = ax.transData.transform_point((x2, y2))
rise = (sp2[1] - sp1[1])
run = (sp2[0] - sp1[0])
slope_degrees = np.degrees(np.arctan2(rise, run))
text.set_rotation_mode('anchor')
text.set_rotation(slope_degrees)
ax.set_ylim(ylim)
return text
@lzkelley
Copy link
Author

_Note_: this method needs to be called after the figure layout is finalized, otherwise things will be altered.

import numpy as np
import matplotlib.pyplot as plt

...
fig, axes = plt.subplots()
color = 'blue'
line, = axes.plot(xdata, ydata, '--', color=color)
...
label_line(line, "Some Label", x, y, color=color)

@mattpitkin
Copy link

mattpitkin commented Nov 23, 2017

This function is great! Just to suggest an edit though, it would be good to check if the axis is a log axis, e.g., from line 36 change:

elif halign.startswith('c'):        
    xx = 0.5*(x1 + x2)
    halign = 'center'

to:

elif halign.startswith('c'):        
    if ax.get_xscale() == 'log':
        xx = 10**(0.5*(np.log10(x1) + np.log10(x2))
    else:
        xx = 0.5*(x1 + x2)
    halign = 'center'

otherwise the positioning is wrong. And also for line 42 change:

yy = np.interp(xx, xdata, ydata)

to

if ax.get_xscale() == 'log' and ax.get_yscale() == 'log':
    yy = 10**(np.interp(np.log10(xx), np.log10(xdata), np.log10(ydata)))
elif ax.get_xscale() == 'log' and ax.get_yscale() != 'log':
    yy = np.interp(np.log10(xx), np.log10(xdata), ydata)
else:
    yy = np.interp(xx, xdata, ydata)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment