In [1]:
%matplotlib inline
import matplotlib.pyplot as pl
import numpy as np
import seaborn as sb

sb.set_style('darkgrid')
sb.set_context('notebook', font_scale=1.5)
blue, green, red, purple = sb.color_palette("deep",4)

Control flow

In the first two notebooks, we went over data types, operators, methods, functions, modules, and packages. Control flow provides structure to programs, enabling more complex manipulation of data inside functions, modules, and packages.

In this section, we'll be talking about:

  • if, elif, and else
  • for loops and comprehensions
  • while loops
  • (and a few other things).

if, elif, and else

We can use if, elif, and else to test if certain conditions hold, executing different blocks of code, depending on the outcome of the test(s). Note that elif is a contraction of else if. The basic syntax is as follows:

if statement:
    # code to execute if statement evaluates as True
elif another_statement:
    # code to execute if statement evaluated as False and
    # another_statement evaluates as True
elif yet_another_statement:
    # code to execute if both statement and another_statement evaluated as False
    # and yet_another_statement evaluates as True
else:
    # code to execute if statement, another_statement, and yet_another_statement
    # all evaluated as False

We can illustrate if and else with an example function that determines if an input number is even or odd.

First, though, remember the modulus operator %, which gives the remainder of the number on the left divided by the number on the right, e.g.,

In [2]:
13%5
Out[2]:
3
In [3]:
14%5
Out[3]:
4
In [4]:
15%5
Out[4]:
0

We can use 2 on the right side of % with a number x to test if x is even or odd. Putting it in a function, we could do this:

In [5]:
def even_or_odd(num):
    mod = num % 2 # divide by two, get remainder
    if mod == 0: # if no remainder...
        print(str(num) + ' is even.')
    else: # yes remainder...
        print(str(num) + ' is odd.')

The first line defines the function. The second line (the first indented line) creates a variable mod equal to the remainder of the input number divided by 2. Then the if statement tests to see if mod is equal to zero. If it is, it tells us that the input is even. Otherwise, it tells us that the input was odd.

Note the use of a new function str(), which converts numbers to strings.

In [6]:
even_or_odd(4904)
4904 is even.
In [7]:
even_or_odd(537)
537 is odd.

As illustrated schematically above, if there are more than two possible test results, we can put elif statements between the initial if and the final else.

Here's a function that tests if an input number is negative, equal to zero, or positive:

In [8]:
def number_type(num):
    if num < 0:
        print(str(num) + ' is negative.')
    elif num == 0:
        print(str(num) + ' is zero.')
    else:
        print(str(num) + ' is positive.')
In [9]:
number_type(-33)
-33 is negative.
In [10]:
number_type(537)
537 is positive.
In [11]:
number_type(0)
0 is zero.

Recall that Booleans can be mapped to 0 (False) or 1 (True):

In [12]:
number_type(False)
False is zero.

When using if, neither elif nor else are required. That is, you can have an if statement by itself. In this case, if the expression evaluates as True, the indented code below the statement will execute, whereas if the expression evaluates as False, it won't.

for loops

We can also have Python iterate over any iterable variable type, executing code at each step. An iterable is an object that is capable of returning its elements one at a time. We've already seen a few iterable data types, e.g., strings, lists, tuples, and arrays.

Here's a simple for loop that iterates through an array of numbers, calling our function number_type() for each one:

In [13]:
# 10 random normal variates with mean 0 and SD 1
nrv = np.random.normal(size=10)

# print array
print(nrv)

# the variable v takes each value in nrv in turn
for v in nrv:
    print(v)
    number_type(np.round(v,3))
    print(' ')
[ 0.00386103 -1.62633267 -1.52297232  0.40286269 -1.06052361 -1.3015004
 -1.51686802 -1.32246949 -0.24609612 -0.28826602]
0.0038610294347037294
0.004 is positive.
 
-1.626332671271649
-1.626 is negative.
 
-1.5229723220184506
-1.523 is negative.
 
0.40286268954277415
0.403 is positive.
 
-1.0605236126225726
-1.061 is negative.
 
-1.301500403571717
-1.302 is negative.
 
-1.5168680191506525
-1.517 is negative.
 
-1.322469485016741
-1.322 is negative.
 
-0.24609612239248818
-0.246 is negative.
 
-0.2882660242338909
-0.288 is negative.
 

Here's a function that tests if a character is a vowel or consonant (note the use of another new operator, in):

In [14]:
def con_or_vow(char_in):
    char = char_in.lower() # make everything lower case
    vowels = list('aeiou') # make the string into a list
    if char in vowels: # test if char is a vowel
        print(char_in + ' is a vowel.')
    elif char == 'y': # test if char is 'y'
        print(char_in + ' might be a vowel.')
    else: # if not a vowel or y, char must be a consonant
        print(char_in + ' is a consonant.')

Using in to check if a variable is an element (or subset) of another variable:

In [15]:
'a' in 'abcde'
Out[15]:
True

in works with lists, too:

In [16]:
list('aeiou')
Out[16]:
['a', 'e', 'i', 'o', 'u']
In [17]:
'a' in list('aeiou')
Out[17]:
True
In [18]:
'ae' in list('aeiou')
Out[18]:
False

Here's the function con_or_vow() in practice:

In [19]:
con_or_vow('a')
a is a vowel.
In [20]:
con_or_vow('b')
b is a consonant.
In [21]:
con_or_vow('y')
y might be a vowel.
In [22]:
con_or_vow('E')
E is a vowel.
In [23]:
con_or_vow('aei')
aei is a consonant.

Here is a for loop that calls con_or_vow() for each character in a string, using the string as an iterable to define the for loop. Note the erroneous responses that this toy function produces when given a space as an input.

In [24]:
question = 'What is a vowel?'

for character in question:
    con_or_vow(character)
W is a consonant.
h is a consonant.
a is a vowel.
t is a consonant.
  is a consonant.
i is a vowel.
s is a consonant.
  is a consonant.
a is a vowel.
  is a consonant.
v is a consonant.
o is a vowel.
w is a consonant.
e is a vowel.
l is a consonant.
? is a consonant.

Base Python has the range() function, and NumPy has the arange() function, both of which are useful if you need to iterate over a sequential set of numbers. (In Python 2, range() returned a list, but in Python 3 the function range() returns a range object, which, along with lists and tuples, is a basic sequence type).

Here is what a range sequence type looks like on its own:

In [25]:
range(10)
Out[25]:
range(0, 10)

These are very useful when creating for loops. Here's a bit of code illustrating the use of range() to create a complex sinusoid:

In [26]:
#np.random.seed(54322) # set seed to fix random number generation
N = 5 # N components
t = np.linspace(0,1,10000) # 1 second time vector
s = np.zeros(len(t))

fig, axes = pl.subplots(3, 1, figsize=(15,15), sharey=True)
ax1, ax2, ax3 = axes

ll = []

for k in range(N):
    f = 2*k + 1 # frequency
    A = 3*np.random.random() # random amplitude
    phi = 2*np.pi*np.random.random() # random phase
    c = A*np.cos(2*np.pi*f*t + phi) # new sinusoid
    ax1.plot(t, c) # plot each individual sinusoid
    s = s + c # add to old sinusoid
    lt, = ax2.plot(t,s) # plot s after adding each sinusoid c
    ll.append(lt)

ax2.legend(ll,['step ' + str(k) for k in range(N)],loc=1)
ax3.plot(t, s, 'k-', lw=3); # plot the final complex sinusoid

You can nest for loops inside for loops (inside for loops, etc...).

Python also has the very useful function enumerate(), which takes an iterable as input and returns tuples of indices and elements of the iterable. Before illustrating the use of enumerate(), though, it's useful to understand the concept of tuple unpacking. The basic idea is that you can assign the elements of a tuple to multiple individual variables in a single line of code:

In [27]:
# example tuple
ex_tup = ('hello', 45232, {'x':4.9e3})
# unpacking
s, n, d = ex_tup
print(s)
print(n)
print(d)
hello
45232
{'x': 4900.0}

One upshot of this concept is that, if a function returns a tuple, you can assign the tuple to a single variable, or you can unpack it and assign the elements to multiple variables.

On each step of a for loop, enumerate() returns a tuple (index, element) consisting of the current index and the element at that index in the iterable being enumerated.

Here's an illustration of nested for loops with enumerate() and some new plotting functionality (note that this example can be done more simply without the for loops):

In [28]:
from mpl_toolkits.mplot3d import Axes3D # for making 3D plot
from matplotlib import cm # for determining colors in plot

fig = pl.figure(figsize=(10,10))
ax = fig.gca(projection='3d')

xv = np.linspace(-3,3,100)
yv = np.linspace(-4,4,100)
X, Y = np.meshgrid(xv,yv) # 2D arrays with xv in each row of X, yv in each column of Y
Z = np.zeros((len(yv),len(xv)))

for xi, x in enumerate(xv): # loop through xv
    for yi, y in enumerate(yv): # loop through yv
        # use xi, yi to index Z, use x, y to calculate 2D bell curve values
        Z[yi,xi] = np.exp( -(x**2 + y**2)/2 )

ax.plot_surface(X, Y, Z, cmap=cm.jet)
ax.view_init(elev=25, azim=35) # adjust elevation and azimuth of plot POV
In [29]:
X
Out[29]:
array([[-3.        , -2.93939394, -2.87878788, ...,  2.87878788,
         2.93939394,  3.        ],
       [-3.        , -2.93939394, -2.87878788, ...,  2.87878788,
         2.93939394,  3.        ],
       [-3.        , -2.93939394, -2.87878788, ...,  2.87878788,
         2.93939394,  3.        ],
       ...,
       [-3.        , -2.93939394, -2.87878788, ...,  2.87878788,
         2.93939394,  3.        ],
       [-3.        , -2.93939394, -2.87878788, ...,  2.87878788,
         2.93939394,  3.        ],
       [-3.        , -2.93939394, -2.87878788, ...,  2.87878788,
         2.93939394,  3.        ]])
In [30]:
Y
Out[30]:
array([[-4.        , -4.        , -4.        , ..., -4.        ,
        -4.        , -4.        ],
       [-3.91919192, -3.91919192, -3.91919192, ..., -3.91919192,
        -3.91919192, -3.91919192],
       [-3.83838384, -3.83838384, -3.83838384, ..., -3.83838384,
        -3.83838384, -3.83838384],
       ...,
       [ 3.83838384,  3.83838384,  3.83838384, ...,  3.83838384,
         3.83838384,  3.83838384],
       [ 3.91919192,  3.91919192,  3.91919192, ...,  3.91919192,
         3.91919192,  3.91919192],
       [ 4.        ,  4.        ,  4.        , ...,  4.        ,
         4.        ,  4.        ]])

Comprehensions

Python also has a for-loop-ish method for creating lists called a "list comprehension." (We'll discuss tuple comprehensions and dictionary comprehensions below.) A list comprehension looks like this:

In [31]:
L = [i**2 for i in range(10)] # squared values for 0, 1, ..., 9
L
Out[31]:
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

As usual, the elements of the list can be anything, e.g.,

In [32]:
L = [(i,i**2, str(i)) for i in range(10)]
L
Out[32]:
[(0, 0, '0'),
 (1, 1, '1'),
 (2, 4, '2'),
 (3, 9, '3'),
 (4, 16, '4'),
 (5, 25, '5'),
 (6, 36, '6'),
 (7, 49, '7'),
 (8, 64, '8'),
 (9, 81, '9')]
In [33]:
L = [str(i) + ' is my favorite number' for i in range(3)]
L
Out[33]:
['0 is my favorite number',
 '1 is my favorite number',
 '2 is my favorite number']

List comprehensions can have "nested for loops", with the former for ... in ... bit functioning as the outer loop, and the latter for ... in ... bit as the inner loop:

In [34]:
L = [(x,y) for x in range(0,5,2) for y in range(0,12,3)]
L
Out[34]:
[(0, 0),
 (0, 3),
 (0, 6),
 (0, 9),
 (2, 0),
 (2, 3),
 (2, 6),
 (2, 9),
 (4, 0),
 (4, 3),
 (4, 6),
 (4, 9)]

We can build additional structure, making a "2D" list, with nested comprehensions:

In [35]:
L = [[(x,y) for x in range(0,5,2)] for y in range(0,12,3)]
L
Out[35]:
[[(0, 0), (2, 0), (4, 0)],
 [(0, 3), (2, 3), (4, 3)],
 [(0, 6), (2, 6), (4, 6)],
 [(0, 9), (2, 9), (4, 9)]]

List comprehensions can also use if and else statements, though they look a little different than the if and else statements discussed above:

In [36]:
L = [np.round(x) if x > 0 else 'NEGATIVE!' for x in np.random.normal(loc=0,scale=5,size=10)]
L
Out[36]:
['NEGATIVE!',
 'NEGATIVE!',
 'NEGATIVE!',
 'NEGATIVE!',
 'NEGATIVE!',
 7.0,
 'NEGATIVE!',
 1.0,
 'NEGATIVE!',
 'NEGATIVE!']

If you just want an if condition (with no corresponding else), you write it like this:

In [37]:
L = [np.round(x,3) for x in np.random.normal(loc=0,scale=5,size=10) if x > 0]
L
Out[37]:
[0.757, 3.369, 2.162, 0.524, 2.457, 13.75]

Here is a toy example of a dictionary comprehension:

In [38]:
keys = ['a','b','c','d']
vals = [0,1,2,3]
D = {k:v for k,v in zip(keys,vals)}
D
Out[38]:
{'a': 0, 'b': 1, 'c': 2, 'd': 3}

And here is a tuple comprehension. Note that the comprehension creates a generator, which we can use in a for loop (rather than being able to look at the results directly, as we did with list and dictionary comprehensions). We can also use the function list() to turn a generator into a list:

In [39]:
T = (x**3 for x in np.arange(-4,5))
T
Out[39]:
<generator object <genexpr> at 0x1a0e613780>
In [40]:
list(T)
Out[40]:
[-64, -27, -8, -1, 0, 1, 8, 27, 64]

Once a generator has returned all of its elements, it's used up:

In [41]:
T = (x**3 for x in np.arange(-4,5))
for t in T:
    print(t)
list(T)
-64
-27
-8
-1
0
1
8
27
64
Out[41]:
[]

while loops

Whereas for loops iterate through a known number of items, while loops iterate until a certain condition is met. More specifically, the code in a while loop will execute repeatedly as long as the specified condition is True.

Here's a simple example that will generate a random number on each iteration, and as long as the random number it generates is less than a threshold we set, it will continue. As soon as the random number exceeds the threshold, it will stop.

In [42]:
th = .75 # threshold
rv = .1 # initialize random variable

while rv < th:
    rv = np.random.random()
    print(rv, rv < th)
    
print("All done!")
0.735741744762914 True
0.10789946411407447 True
0.7382071184822794 True
0.838886115538929 False
All done!

Here is a more complicated example of the same thing. In this case, we keep track of how many times we've generated a random number, and the while loop is exited if it goes on for too long:

In [43]:
th = .99 # very high threshold
rv = .5 # initial "random" value

cont = rv < th # condition for while loop

n_steps = 0 # initial number of steps taken

while cont:
    rv = np.random.random()
    cont = rv < th # update condition
    n_steps += 1 # add 1 to number of steps; n_steps = n_steps + 1
    if n_steps > 10: # test to see if number of steps is greater than 10
        print("Breaking the loop. Too many steps.")
        break # exit while loop
    print(rv,cont,n_steps)
    
print("All done!")
0.20892688121938396 True 1
0.44550145224465976 True 2
0.5825055748916481 True 3
0.5928755869155782 True 4
0.9584358713018646 True 5
0.5892216811298188 True 6
0.8140151768993572 True 7
0.18193621576053498 True 8
0.7748001543897196 True 9
0.6454247775740259 True 10
Breaking the loop. Too many steps.
All done!

A more realistic example

Here's a more realistic, and complex, example illustrating a while loop that uses the Newton-Raphson algorithm to find the root (zero) of a function. This method uses calculus to generate approximations to the root of a function, and we can use a while loop to execute the algorithm until the change from one approximation to the next is smaller than some threshold that we choose.

We'll use it to find the root of the polynomial $f(x) = a + bx + cx^3$.

To use the Newton-Raphson methods, we need the derivative of this function, which is $f^\prime(x) = \displaystyle{\frac{dy}{dx}} = b + 3cx^2$.

Let $x_0$ be our initial guess for the root of the function. The approximation at the first step is $x_1 = x_0 - \displaystyle{\frac{f(x_0)}{f^\prime(x_0)}}$.

The approximation at step $n+1$ is $x_{n+1} = x_n - \displaystyle{\frac{f(x_n)}{f^\prime(x_n)}}$.

We can run the algorithm until the approximation at step $n$ is less than a certain distance from the approximation at step $n-1$.

First, we'll define some functions for evaluating the function and its derivative and the algorithm at any given step:

In [44]:
# the original function
def f_of_x(x,a=2,b=5,c=3):
    return a + b*x + c*x**3

# the derivative
def d_f_of_x(x,b=5,c=3):
    return b + 3*c*x**2

# Newton-Raphson algorithm
def newtrap(x,f,g,a=2,b=5,c=3):
    return x - f(x,a,b,c)/g(x,b,c)

We'll plot the function and its derivative for a specific set of coefficients $a, b, c$ to see what we're dealing with:

In [45]:
a, b, c = -10, 15, -5
xv = np.linspace(-3,3,500)
f_x = f_of_x(xv,a,b,c)
g_x = d_f_of_x(xv,b,c)
fig, ax = pl.subplots(1, 1, figsize=(12,6))
ax.plot(xv,np.zeros(len(xv)),':',color=[.5,.5,.5,.5])
ax.plot(xv,f_x, lw=3)
ax.plot(xv,g_x,'--', lw=3);
In [46]:
xo = 3 # initial guess
xn = newtrap(xo, f_of_x, d_f_of_x, a=a, b=b, c=c) # next algorithm step
th = .00001 # threshold for estimate change from step to step
d = np.abs(xn-xo) # distance between estimates at steps 0 and 1
n_steps = 1

approx = [xo,xn]

while d > th: # as long as xo and xn are sufficiently far apart...
    xo = xn # use xn as new xo
    xn = newtrap(xo, f_of_x, d_f_of_x, a=a, b=b, c=c) # get next algorithm estimate
    d = np.abs(xn-xo) # recalculate distance between old and new estimates
    approx.append(xn) # append new estimate to list of estimates
    n_steps += 1 # add to number of steps taken
    
approx = np.array(approx)
    
print(xn,n_steps)
1.000006024834514 19
In [47]:
fig, ax = pl.subplots(1, 1, figsize=(12,6))
line, = ax.plot([],[],'r') #
xv = np.linspace(-3,3,500)
f_x = f_of_x(xv,a,b,c)
ax.plot(xv,np.zeros(len(xv)),':',color=[.5,.5,.5,.5], lw=3)
ax.plot(xv,f_x,'b--')
for xi,yi in zip(approx,f_of_x(approx,a,b,c)):
    ax.plot(xi,yi,'o',ms=10)
for ii,xyi in enumerate(zip(approx,f_of_x(approx,a,b,c))):
    xi,yi = xyi
    ax.text(xi-.1,yi+.5,str(ii),fontsize=12)
In [48]:
approx # estimates of the root
Out[48]:
array([3.        , 2.16666667, 1.65497076, 1.35441516, 1.18609939,
       1.09569009, 1.04857325, 1.02447858, 1.01228862, 1.00615682,
       1.00308156, 1.00154157, 1.00077098, 1.00038554, 1.00019278,
       1.00009639, 1.0000482 , 1.0000241 , 1.00001205, 1.00000602])
In [49]:
f_of_x(approx, a, b, c) # function f(x) at estimates approach zero...
Out[49]:
array([-1.00000000e+02, -2.83564815e+01, -7.83966917e+00, -2.10674223e+00,
       -5.51720649e-01, -1.41729878e-01, -3.59634238e-02, -9.06134992e-03,
       -2.27443083e-03, -5.69762801e-04, -1.42586266e-04, -3.56648368e-05,
       -8.91849767e-06, -2.22991077e-06, -5.57513506e-07, -1.39382855e-07,
       -3.48462716e-08, -8.71163941e-09, -2.17791829e-09, -5.44480017e-10])

We can also implement a version of the Newton-Raphson algorithm using a for loop, in which case we will have to decide on a number of iterations ahead of time, and we won't stop according to a criterion based on how good our approximation is (i.e., we will illustrate why a while loop is better for this situation than a for loop).

In [50]:
xo = 3
xn = newtrap(xo, f=f_of_x, g=d_f_of_x, a=a, b=b, c=c)
n_iter = 2

approx_for = [xo,xn]

for i in range(n_iter):
    xo = xn
    xn = newtrap(xo, f=f_of_x, g=d_f_of_x, a=a, b=b, c=c)
    approx_for.append(xn)
    
approx_for = np.array(approx_for)
    
fig, ax = pl.subplots(1, 1, figsize=(12,6))
line, = ax.plot([],[],'r') #
xv = np.linspace(-3,3,500)
f_x = f_of_x(xv,a,b,c)
ax.plot(xv,np.zeros(len(xv)),':',color=[.5,.5,.5,.5], lw=3)
ax.plot(xv,f_x,'b--')
for xi,yi in zip(approx_for,f_of_x(approx_for,a,b,c)):
    ax.plot(xi,yi,'o',ms=10)
for ii,xyi in enumerate(zip(approx_for,f_of_x(approx_for,a,b,c))):
    xi,yi = xyi
    ax.text(xi-.1,yi+.5,str(ii),fontsize=12)

Warning!

You have to be careful with while loops, since you can get stuck if the criterion governing whether or not it continues never changes. I will illustrate in IPython, since there does not seem to be a simple way to escape from an infinite loop in a notebook other than to close the tab/browser. In IPython, you can type ctrl-c to force it to quit.