How to draw curves

5.1 Introduction

This how-to will demonstrate how to draw mathematical curves in Python.

Imagine that we want to plot three segments of the function:

$f(x) = x(x-5)$

That can be done by selecting $x$ -values in the three segments and evaluate the corresponding $y$ -values as done here:

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

plt.rc('font', size=18)
fig, ax = plt.subplots()

x = np.linspace(-5,0,100)
y = x * (x - 5)
ax.plot(x,y,linewidth=3,color='royalblue')

x = np.linspace(0,5,100)
y = x * (x - 5)
ax.plot(x,y,linewidth=3,color='crimson')

x = np.linspace(5,10,100)
y = x * (x - 5)
ax.plot(x,y,linewidth=3,color='teal')

ax.grid()
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim([-5,10])
ax.legend(['$x<0$','$0<x<5$','$x>5$'])

Figure 1

Notice how we end up writing the exact same code, y = x * (x - 5), for the function three times as we have new $x$ -values. This is undesirable as it invites to making coding mistakes. Imagine for instance that the function is changed. Then the code must be changed consistently at three different places.

5.2 Advanced loop

One solution that makes the function expression appear only once in the code woud be to introduce a loop. Since more than one thing is changing within the loop, it needs to be slightly more complicated than loops we have seen previously. Here is an example:

fig, ax = plt.subplots()

for x1, x2, col in [(-5, 0, 'deepskyblue'),
                    (0, 5, 'coral'),
                    (5, 10, 'turquoise')]:
    x = np.linspace(x1,x2,100)
    y = x * (x - 5)
    ax.plot(x,y,linewidth=3,color=col)

ax.grid()
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim([-5,10])
ax.legend(['$x<0$','$0<x<5$','$x>5$'])

Figure 2

where the list contain elements that are tuples. When the loop is executed the tuples are dealt with one by one. First the (-5, 0, 'deepskyblue') tuple is handled, and the three variables x1, x2, col specified between for and in assume the three values from the tuple, i.e. -5, 0, and 'deepskyblue', respectively. Then follows the next tuple, and so on.

5.3 Even more advanced loop

If the desired boundaries for the plotted segments and the colors are provided in three different lists (rather than as tuples assembled in a list) one may invoke the zip function to create the tuples as in the previous example. It takes the form:

fig, ax = plt.subplots()

x1_vals = [-5, 0, 5]
x2_vals = [0, 5, 10]
colors = ['deepskyblue', 'coral', 'turquoise']

for x1, x2, col in zip(x1_vals, x2_vals, colors):
    x = np.linspace(x1,x2,100)
    y = x * (x - 5)
    ax.plot(x,y,linewidth=3,color=col)

ax.grid()
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim([-5,10])
ax.legend(['$x<0$','$0<x<5$','$x>5$'])

Figure 3

where the first tuple from zip(x1_vals, x2_vals, colors) becomes exactly (-5, 0, 'deepskyblue'), the next becomes (0, 5, 'coral'), and so on

5.4 Lambda function

Another way to circumvent that the function expression is typed in several times in the code is to actually define the function. We have already worked with e.g. the cosine function, which became available as np.cos when the NumPy package was imported as np. When writing a mathematical function in Python, the shortest syntax possible involves a so-called lambda function. Here is an example:

f = lambda x: x * (x - 5)
f(1), f(np.array([1, 2]))

(-4, array([-4, -6]))

It is seen, that a lambda function may be called with a scalar value or with a NumPy array of values.

Now, the plotting of segments of a curve as considered above can be written in this way:

fig, axs = plt.subplots()

f = lambda x: x * (x - 5)

x = np.linspace(-5,0,100)
axs.plot(x,f(x),linewidth=3,color='cornflowerblue')

x = np.linspace(0,5,100)
axs.plot(x,f(x),linewidth=3,color='tomato')

x = np.linspace(5,10,100)
axs.plot(x,f(x),linewidth=3,color='olive')

axs.grid()
axs.set_xlabel('$x$')
axs.set_ylabel('$y$')
axs.set_xlim([-5,10])
axs.legend(['$x<0$','$0<x<5$','$x>5$'])

Figure 4

5.5 Several lambda functions

Imagine that the function we were to plot had the form:

$f_{x_0}(x) = x(x-x_0)$

and that some different values were to be used for $x_0$ . One would then be able to make the plot by introducing three different lambda functions:

fig, axs = plt.subplots()

x = np.linspace(0,10,100)

f1 = lambda x: x * (x - 2)
f2 = lambda x: x * (x - 5)
f3 = lambda x: x * (x - 7)

axs.plot(x,f1(x),linewidth=3,color='royalblue')
axs.plot(x,f2(x),linewidth=3,color='crimson')
axs.plot(x,f3(x),linewidth=3,color='teal')

axs.grid()
axs.set_xlabel('$x$')
axs.set_ylabel('$y$')
axs.set_xlim([0,10])
axs.legend(['$x_0=2$','$x_0=5$','$x_0=7$'])

Figure 5

5.6 Lambda function returning lambda function

In the previous example we had three functions of the same form differing only by the value of $x_0$ . A neat way to write these three functions in the Python code would therefor be to write a lambda function that takes $x_0$ as and returns the function for that particular value of $x_0$ . That would look like this:

fig, axs = plt.subplots()

x = np.linspace(0,10,100)

f = lambda x0: lambda x: x * (x - x0)

f1 = f(3)
f2 = f(6)
f3 = f(8)

axs.plot(x,f1(x),linewidth=3,color='steelblue')
axs.plot(x,f2(x),linewidth=3,color='firebrick')
axs.plot(x,f3(x),linewidth=3,color='seagreen')

axs.grid()
axs.set_xlabel('$x$')
axs.set_ylabel('$y$')
axs.set_xlim([0,10])
axs.legend(['$x_0=3$','$x_0=6$','$x_0=8$'])

Figure 6

You might think this looks a little bit complicated if this is the first time you use a function. That is ok. Then the solution model with three separately written lambda functions is fine.

5.7 Curves with separate axes

It is possible to have more than one set of axes and thereby plot curves separately. It is done by providing arguments in the subplots call. fig, axs = subplot(2,1) for instance provides two rows and one coloumn of axes, i.e. two axes atop each other. The second variable, axs, in the assignment statement now becomes a list of Axes object instances. They can then be addressed one by one, e.g. by invoking ax = axs[0] and ax = axs[1] as in this example:

fig, axs = plt.subplots(2,1,sharex=True,figsize=(6,7.5))
fig.suptitle('Two Gaussians')

f_of_x0 = lambda x0: lambda x: np.exp(-np.power(x-x0,2))
xs = np.linspace(-5,5,100)
x0s = [-1, 0]

ax = axs[0]
x0 = x0s[0]
f = f_of_x0(x0)
ax.plot(xs,f(xs))
ax.grid()
ax.set_ylabel('$y$')
ax.set_title('$x_0={}$'.format(x0))

ax = axs[1]
x0 = x0s[1]
f = f_of_x0(x0)
ax.plot(xs,f(xs))
ax.grid()
ax.set_ylabel('$y$')
ax.set_title('$x_0={}$'.format(x0))

ax.set_xlabel('$x$')
fig.tight_layout()
fig.subplots_adjust(top=0.9)

Figure 7

5.8 Curves with separate axes II

More compact and more efficient coding results if a loop is used to run through the axes, as in this example:

fig, axs = plt.subplots(4,1,sharex=True,figsize=(6,10))
fig.suptitle('Moving Gaussians')

f_of_x0 = lambda x0: lambda x: np.exp(-np.power(x-x0,2))

x0s = [-1, 0, 1, 2]

xs = np.linspace(-5,5,100)
for ax,x0 in zip(axs,x0s):
    f = f_of_x0(x0)
    ax.plot(xs,f(xs))
    ax.grid()
    ax.set_ylabel('$y$')
    ax.set_title('$x_0={}$'.format(x0))
    
ax.set_xlabel('$x$')
fig.tight_layout()
fig.subplots_adjust(top=0.9)

Figure 8