Search code examples
pythonmatplotlibline-plotplot-annotationslabellines

How to place inline labels in a line plot


In Matplotlib, it's not too tough to make a legend (example_legend(), below), but I think it's better style to put labels right on the curves being plotted (as in example_inline(), below). This can be very fiddly, because I have to specify coordinates by hand, and, if I re-format the plot, I probably have to reposition the labels. Is there a way to automatically generate labels on curves in Matplotlib? Bonus points for being able to orient the text at an angle corresponding to the angle of the curve.

import numpy as np
import matplotlib.pyplot as plt

def example_legend():
    plt.clf()
    x = np.linspace(0, 1, 101)
    y1 = np.sin(x * np.pi / 2)
    y2 = np.cos(x * np.pi / 2)
    plt.plot(x, y1, label='sin')
    plt.plot(x, y2, label='cos')
    plt.legend()

Figure with legend

def example_inline():
    plt.clf()
    x = np.linspace(0, 1, 101)
    y1 = np.sin(x * np.pi / 2)
    y2 = np.cos(x * np.pi / 2)
    plt.plot(x, y1, label='sin')
    plt.plot(x, y2, label='cos')
    plt.text(0.08, 0.2, 'sin')
    plt.text(0.9, 0.2, 'cos')

Figure with inline labels


Solution

  • Nice question, a while ago I've experimented a bit with this, but haven't used it a lot because it's still not bulletproof. I divided the plot area into a 32x32 grid and calculated a 'potential field' for the best position of a label for each line according the following rules:

    • white space is a good place for a label
    • Label should be near corresponding line
    • Label should be away from the other lines

    The code was something like this:

    import matplotlib.pyplot as plt
    import numpy as np
    from scipy import ndimage
    
    
    def my_legend(axis = None):
    
        if axis == None:
            axis = plt.gca()
    
        N = 32
        Nlines = len(axis.lines)
        print Nlines
    
        xmin, xmax = axis.get_xlim()
        ymin, ymax = axis.get_ylim()
    
        # the 'point of presence' matrix
        pop = np.zeros((Nlines, N, N), dtype=np.float)    
    
        for l in range(Nlines):
            # get xy data and scale it to the NxN squares
            xy = axis.lines[l].get_xydata()
            xy = (xy - [xmin,ymin]) / ([xmax-xmin, ymax-ymin]) * N
            xy = xy.astype(np.int32)
            # mask stuff outside plot        
            mask = (xy[:,0] >= 0) & (xy[:,0] < N) & (xy[:,1] >= 0) & (xy[:,1] < N)
            xy = xy[mask]
            # add to pop
            for p in xy:
                pop[l][tuple(p)] = 1.0
    
        # find whitespace, nice place for labels
        ws = 1.0 - (np.sum(pop, axis=0) > 0) * 1.0 
        # don't use the borders
        ws[:,0]   = 0
        ws[:,N-1] = 0
        ws[0,:]   = 0  
        ws[N-1,:] = 0  
    
        # blur the pop's
        for l in range(Nlines):
            pop[l] = ndimage.gaussian_filter(pop[l], sigma=N/5)
    
        for l in range(Nlines):
            # positive weights for current line, negative weight for others....
            w = -0.3 * np.ones(Nlines, dtype=np.float)
            w[l] = 0.5
    
            # calculate a field         
            p = ws + np.sum(w[:, np.newaxis, np.newaxis] * pop, axis=0)
            plt.figure()
            plt.imshow(p, interpolation='nearest')
            plt.title(axis.lines[l].get_label())
    
            pos = np.argmax(p)  # note, argmax flattens the array first 
            best_x, best_y =  (pos / N, pos % N) 
            x = xmin + (xmax-xmin) * best_x / N       
            y = ymin + (ymax-ymin) * best_y / N       
    
    
            axis.text(x, y, axis.lines[l].get_label(), 
                      horizontalalignment='center',
                      verticalalignment='center')
    
    
    plt.close('all')
    
    x = np.linspace(0, 1, 101)
    y1 = np.sin(x * np.pi / 2)
    y2 = np.cos(x * np.pi / 2)
    y3 = x * x
    plt.plot(x, y1, 'b', label='blue')
    plt.plot(x, y2, 'r', label='red')
    plt.plot(x, y3, 'g', label='green')
    my_legend()
    plt.show()
    

    And the resulting plot: enter image description here