Search code examples
pythontextmatplotlibcurve

Matplotlib, adding text with more than one line. Adding text that can follow the curve


I have added text to a plot, coded in each line, and then adjusted it look decent, increase or decrease the width, or change the placement. However, is there a way to have Python know where you want the text and how you want it set? Then I could add the text and Python would work out the details.

For example, take a look at the image below:

enter image description here

In the figure, I have 3 lines of text in the upper left corner and one line above the line of the plot.

I had to adjust the 3 lines to get a decent spacing. This wasnt a difficult task but it would be easy if I could say here is the text, here is the location, and then Python stacks it with proper spacing.

For the lone line, I had to make adjustments so it wasn't on the line and lower the line. For this case, is is possible to tell python I would like the text above the plot and 80% down the line?

I am used to LaTeX where I can make this adjustments without hard coding the coordinates. The advantage are

(1) if I want to change the location, I can change the percentage shift and not the coordinate.

(2) if the line is angled, the text will adjust to the line.

The advantage to (2) is that I am trying to put text on the top portion of the figure that slopes upward with the line.

Can this be done or am I asking to much? If so, how do I do this?

Here is the code that implements the figure:

import numpy as np
import pylab

r1 = 1  #  AU Earth                                                                 
r2 = 1.524  #  AU Mars                                                              
deltanu = 75 * np.pi / 180  #  angle in radians                                     
mu = 38.86984154054163

c = np.sqrt(r1 ** 2 + r2 ** 2 - 2 * r1 * r2 * np.cos(deltanu))   
s = (r1 + r2 + c) / 2
am = s / 2


def g(a):
    alphag = 2* np.pi - 2 * np.arcsin(np.sqrt(s / (2 * a)))
    return (np.sqrt(a ** 3 / mu)
            * (alphag - betag - (np.sin(alphag) - np.sin(betag))))


def f(a):
    alpha = 2 * np.arcsin(np.sqrt(s / (2 * a)))
    beta = 2 * np.arcsin(np.sqrt((s - c) / (2 * a)))
    return (np.sqrt(a **3 / mu) * (alpha - betag - (np.sin(alpha)
                                                      - np.sin(betag))))


betag = -2 * np.arcsin(np.sqrt((s - c) / (2 * a)))
a = np.linspace(am, 2, 500000)

a = np.linspace(am, 2, 500000)

fig = pylab.figure()
ax = fig.add_subplot(111)
ax.plot(a, f(a), color = '#000000')
ax.plot(a, g(a), color = '#000000')
pylab.xlim((0.9, 2))
pylab.ylim((0, 2))

pylab.xlabel('Semi-major Axis $a$ in AU')
pylab.ylabel('Time of Flight in Years')
pylab.text(1, 1.8, '$r_1 = 1.0$ AU', fontsize = 11, color = 'r')
pylab.text(1, 1.7, '$r_2 = 1.524$ AU', fontsize = 11, color = 'r')
pylab.text(1, 1.6, '$\\Delta \\nu = 75^{\\circ}$', fontsize = 11,
           color = 'r')
pylab.text(1.75, 0.35, '$\\alpha = \\alpha_0$', fontsize = 11,
           color = 'r')
pylab.savefig('lamberttransferties.eps', format = 'eps')
pylab.show()

Solution

  • You can use line separators \n:

     pylab.text(1, 1.5, '$r_1 = 1.0$ AU\n' +\
                        '$r_2 = 1.524$ AU\n' +\
                        '$\\Delta \\nu = 75^{\\circ}$', fontsize = 11, color = 'r')
    

    pylab.text() uses data coordinates by default, but you can use relative positions (0,0) to the lower-left and (1,1) to the upper-right, passing the parameter transform. See this example:

    pylab.text(0.6, 0.75, 'using axis coords', transform=ax.transAxes)
    

    The parameters: verticalalignment and horizontalalignment can also help you tremendously. Suppose you want to place a texts at the very corners:

    pylab.text(1.,1.,'top-right', transform=ax.transAxes,
               horizontalalignment='right', verticalalignment='top')
    
    pylab.text(0.,0.,'bottom-left', transform=ax.transAxes,
               horizontalalignment='left', verticalalignment='bottom')
    

    enter image description here

    To automatically calculate an angle to the text depending on your data you can do the following approach:

    • detect the data closest point
    • find the a sequence near the closest point and fit a curve using this sequence (the example below uses a fourth order curve)
    • calculate the derivative at the point where you want the text placed
    • correct the direvative with ax.get_data_ratio() OBS: not needed if ax.axis('scaled') is used, for example

    This algorithm can be implemented as follows:

    def rtext(line,x,y,s, **kwargs):
        from scipy.optimize import curve_fit
        xdata,ydata = line.get_data()
        dist = np.sqrt((x-xdata)**2 + (y-ydata)**2)
        dmin = dist.min()
        TOL_to_avoid_rotation = 0.3
        if dmin > TOL_to_avoid_rotation:
            r = 0.
        else:
            index = dist.argmin()
            xs = xdata[ [index-2,index-1,index,index+1,index+2] ]
            ys = ydata[ [index-2,index-1,index,index+1,index+2] ]
            def f(x,a0,a1,a2,a3):
                return a0 + a1*x + a2*x**2 + a3*x**3
            popt, pcov = curve_fit(f, xs, ys, p0=(1,1,1,1))
            a0,a1,a2,a3 = popt
            ax = pylab.gca()
            derivative = (a1 + 2*a2*x + 3*a3*x**2)
            derivative /= ax.get_data_ratio()
            r = np.arctan( derivative )
        return pylab.text(x, y, s, rotation=np.rad2deg(r), **kwargs)
    

    The following test example shows how to use it:

    ax = pylab.subplot(111)
    thetas = np.linspace(0,6*np.pi,1000)
    i = np.arange(len(thetas))
    xdata = (1. + (3.-1.)*i/len(thetas))*np.cos(thetas)
    ydata = (1. + (3.-1.)*i/len(thetas))*np.sin(thetas)
    ax.plot(xdata, ydata, color = 'b')
    pylab.xlabel('x')
    pylab.ylabel('y')
    for x, y in zip(xdata,ydata)[::25]:
        rtext(ax.lines[0], x, y,
              '$\\alpha = \\alpha_0$', fontsize = 14, color = 'r',
              horizontalalignment='center', verticalalignment='center')
    

    enter image description here

    Changing verticalalignment='bottom'

    enter image description here