How to use loops

3.1 Introduction

This how-to will demonstrate more plotting functions and gradually introduce the use of variables, lists, and loops. We will start by making a figure that involves a number of colored triangles:

%matplotlib inline
import matplotlib.pyplot as plt
fig, ax = plt.subplots()

ax.fill([0, 1, 0.5],[0, 0, 0.866])
ax.fill([0, 0.5, -0.5],[0, 0.866, 0.866])
ax.fill([0, -0.5, -1],[0, 0.866, 0])
ax.grid()
ax.set_xticks([-1,-0.5,0,0.5,1])
ax.set_yticks([0,0.5,1])

Figure 1

A number of new methods on the Axes object are used: fill takes two lists with $x$ and $y$ values, respectively, and plots a closed polygon with corners at the $(x,y)$ points. grid adds the grid of lines throughout the figure. The lines appear at the tick values, which are set by the set_xticks and set_yticks methods that each take a list as arguments. Notice how one unit along the $x$ -axis and one unit along the $y$ -axis do not span the same distance in the figure. This is remedied by calling set_aspect('equal') on ax:

ax.set_aspect('equal')
fig

Figure 2

3.2 Variables

When plotting above, many of the $x$ and $y$ values used where typed a number of times. Generally, when writing computer scripts and programs, it is good practise to use so-called variables. When using variables in Python, values are put in "storage boxes" in the computer and the variable names act as labels for these boxes. The content of the boxes may be retrieved using the variable names, whereby any recurring values need only be typed in once in the Python code.

Here is an example where the same figure is made, while the values are now reference via variables:

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

x0 = -1
x1 = -0.5
x2 = 0
x3, x4 = 0.5, 1

y0, y1 = 0, 0.866

ax.fill([x2, x4, x3],[y0, y0, y1])
ax.fill([x2, x3, x1],[y0, y1, y1])
ax.fill([x2, x1, x0],[y0, y1, y0])
ax.grid()
ax.set_xticks([x0, x1, x2, x3, x4])
ax.set_yticks([x2, x3, x4])

ax.text(x0, y0,'(x0,y0)',horizontalalignment='left', verticalalignment='top')
ax.text(x1, y1,'(x1,y1)',horizontalalignment='center', verticalalignment='bottom')
ax.text(x2, y0,'(x2,y0)',horizontalalignment='center', verticalalignment='top')
ax.text(x3, y1,'(x3,y1)',horizontalalignment='center', verticalalignment='bottom')
ax.text(x4, y0,'(x2,y0)',horizontalalignment='right', verticalalignment='top')

ax.set_aspect('equal')

Figure 3

The example further demonstrates how some text can be displayed in various ways inside the figure.

3.3 Lists

In the last section, variables where introduced. However, you might already have wondered if that many variables were really necessary and helpful? To avoid introducing too many variable names, it is often useful to use lists and refer to elements within the lists.

You define a list like this:

x = [-1, -0.5, 0, 0.5, 1]
x

[-1, -0.5, 0, 0.5, 1]

and refer to its elements using an index in square parentheses. The first element has index 0:

x[0]

-1

and e.g. the third element has index 2:

x[2]

Using the colon, :, more elements can be extracted and will result in a new list. For instance:

x[2:]

[0, 0.5, 1]

The figure with the triangular polygons may now be constructed like this:

fig, ax = plt.subplots()

x = [-1, -0.5, 0, 0.5, 1]
y = [0, 0.866]

ax.fill([x[2], x[4], x[3]],[y[0], y[0], y[1]])
ax.fill([x[2], x[3], x[1]],[y[0], y[1], y[1]])
ax.fill([x[2], x[1], x[0]],[y[0], y[1], y[0]])
ax.grid()
ax.set_xticks(x)
ax.set_yticks(x[2:])

ax.set_aspect('equal')

Figure 4

3.4 Calculated values

Unhappy about the explicit use of the value $0.866$ for some of the $y$ -coordinates? Well, you are using a computer, so why not let it calculate any relevant numbers rather than you having to type them in? To do so, you will need to import a python package that can do calculus. We shall be using the package NumPy. As a first example, it can be input and the value of $\pi$ can be printed:

import numpy as np
np.pi

3.141592653589793

Having loaded NumPy we can get to do the calculus necessary for plotting the triangles, e.g.:

np.sqrt(3/4)

0.8660254037844386

or:

np.cos(np.pi/3), np.sin(np.pi/3)

(0.5000000000000001, 0.8660254037844386)

The figure with the three triangles may now be coded like this:

fig, ax = plt.subplots()

p0 = [0, 0]
p1 = [np.cos(0), np.sin(0)]
p2 = [np.cos(np.pi/3),np.sin(np.pi/3)]
xs = [p0[0],p1[0],p2[0]]
ys = [p0[1],p1[1],p2[1]]
ax.fill(xs,ys)

p1 = p2
p2 = [np.cos(np.pi*2/3),np.sin(np.pi*2/3)]
xs = [p0[0],p1[0],p2[0]]
ys = [p0[1],p1[1],p2[1]]
ax.fill(xs,ys)

p1 = p2
p2 = [np.cos(np.pi),np.sin(np.pi)]
xs = [p0[0],p1[0],p2[0]]
ys = [p0[1],p1[1],p2[1]]
ax.fill(xs,ys)

ax.grid()

ax.set_xticks([-1,-0.5,0,0.5,1])
ax.set_yticks([0, 0.5, 1])

ax.set_aspect('equal')

Figure 5

Note how the statement, p1 = p2, renames a list as we move from one triangle to the next. This is strictly not necessary but means that the assignments, xs = [p0[0],p1[0],p2[0]] and ys = [p0[1],p1[1],p2[1]], may look the same for all three triangles, which makes it simpler to write the code correctly.

3.5 Loops

In the previous section, we ended up with a small piece of code, where more or less the same section,

p1 = [np.cos(BLA), np.sin(BLA)]
p2 = [np.cos(BLA-BLA),np.sin(BLA-BLA)]
xs = [p0[0],p1[0],p2[0]]
ys = [p0[1],p1[1],p2[1]]
ax.fill(xs,ys)

was repeated three times. Only, BLA and BLA-BLA differed in the three instances. Now, this can be exploited using Pythons loop construction. If you write e.g.:

for t in [0, 2, 8]:
    print('Before printing',t)
    print(t)
    print('After printing',t)

where the indentation of the three print-statements is part of the syntax, then the three print-statements will be executed a number of times, in fact three times. The first time, t will be $0$ , the next time, t will be $2$ and the final time, t will be $8$ as these are the elements of the list given in the for statement. The output becomes:

Before printing 0
0
After printing 0
Before printing 2
2
After printing 2
Before printing 8
8
After printing 8

Utilizing the Pythons for-loop construction, the triangles can be plotted like this:

fig, ax = plt.subplots()

p0 = [0, 0]

for theta in [0, np.pi/3, np.pi*2/3]:
    p1 = [np.cos(theta), np.sin(theta)]
    p2 = [np.cos(theta + np.pi/3),np.sin(theta + np.pi/3)]
    xs = [p0[0],p1[0],p2[0]]
    ys = [p0[1],p1[1],p2[1]]
    ax.fill(xs,ys)

ax.grid()

ax.set_xticks([-1,-0.5,0,0.5,1])
ax.set_yticks([0, 0.5, 1])

ax.set_aspect('equal')

Figure 6

Notice, how the indentation is terminated at ax.grid() and onwards, since these are statements that should only be conducted once and not over and over again for each value in the list.

3.6 Nested loops

You can have loops inside loops:

fig, ax = plt.subplots()


for x0 in [0, 2]:
    p0 = [x0, 0]
    for theta in [0, np.pi/3, np.pi*2/3]:
        p1 = [x0 + np.cos(theta), np.sin(theta)]
        p2 = [x0 + np.cos(theta + np.pi/3),np.sin(theta + np.pi/3)]
        xs = [p0[0],p1[0],p2[0]]
        ys = [p0[1],p1[1],p2[1]]
        ax.fill(xs,ys)

ax.grid()

ax.set_xticks([-1, 0, 1, 2, 3])
ax.set_yticks([-1, 0, 1, 2])

ax.set_aspect('equal')

Figure 7

3.7 NumPy arrays

At this stage, you might want to start making the arithmetic directly on the lists. I.e. you might want to be able to add coordinates of the lower left corner, p0, of a triangle to the other points, p1 and p2, of the triangle. That is not possible with python lists. You have to use another type of lists for that, NumPy arrays. Here is an example with arithmetic on NumPy arrays:

p0 = np.array([2, 0])
p0

array([2, 0])

p1 = np.array([1, 0])
p1

array([1, 0])

p0 + p1

array([3, 0])

Getting back to the example with triangles it can now look like this:

fig, ax = plt.subplots()

for x0 in [0, 2]:
    p0 = np.array([x0, 0])
    for theta in [0, np.pi/3, np.pi*2/3]:
        p1 = p0 + np.array([np.cos(theta), np.sin(theta)])
        p2 = p0 + np.array([np.cos(theta + np.pi/3),np.sin(theta + np.pi/3)])
        xs = [p0[0],p1[0],p2[0]]
        ys = [p0[1],p1[1],p2[1]]
        ax.fill(xs,ys)

ax.grid()

ax.set_xticks([-1, 0, 1, 2, 3])
ax.set_yticks([-1, 0, 1, 2])

ax.set_aspect('equal')

Figure 8