%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, andelseforloops and comprehensionswhileloops- (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.,
13%5
14%5
15%5
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:
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.
even_or_odd(4904)
even_or_odd(537)
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:
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.')
number_type(-33)
number_type(537)
number_type(0)
Recall that Booleans can be mapped to 0 (False) or 1 (True):
number_type(False)
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:
# 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(' ')
Here's a function that tests if a character is a vowel or consonant (note the use of another new operator, in):
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:
'a' in 'abcde'
in works with lists, too:
list('aeiou')
'a' in list('aeiou')
'ae' in list('aeiou')
Here's the function con_or_vow() in practice:
con_or_vow('a')
con_or_vow('b')
con_or_vow('y')
con_or_vow('E')
con_or_vow('aei')
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.
question = 'What is a vowel?'
for character in question:
con_or_vow(character)
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:
range(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:
#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:
# example tuple
ex_tup = ('hello', 45232, {'x':4.9e3})
# unpacking
s, n, d = ex_tup
print(s)
print(n)
print(d)
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):
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
X
Y
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:
L = [i**2 for i in range(10)] # squared values for 0, 1, ..., 9
L
As usual, the elements of the list can be anything, e.g.,
L = [(i,i**2, str(i)) for i in range(10)]
L
L = [str(i) + ' is my favorite number' for i in range(3)]
L
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:
L = [(x,y) for x in range(0,5,2) for y in range(0,12,3)]
L
We can build additional structure, making a "2D" list, with nested comprehensions:
L = [[(x,y) for x in range(0,5,2)] for y in range(0,12,3)]
L
List comprehensions can also use if and else statements, though they look a little different than the if and else statements discussed above:
L = [np.round(x) if x > 0 else 'NEGATIVE!' for x in np.random.normal(loc=0,scale=5,size=10)]
L
If you just want an if condition (with no corresponding else), you write it like this:
L = [np.round(x,3) for x in np.random.normal(loc=0,scale=5,size=10) if x > 0]
L
Here is a toy example of a dictionary comprehension:
keys = ['a','b','c','d']
vals = [0,1,2,3]
D = {k:v for k,v in zip(keys,vals)}
D
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:
T = (x**3 for x in np.arange(-4,5))
T
list(T)
Once a generator has returned all of its elements, it's used up:
T = (x**3 for x in np.arange(-4,5))
for t in T:
print(t)
list(T)
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.
th = .75 # threshold
rv = .1 # initialize random variable
while rv < th:
rv = np.random.random()
print(rv, rv < th)
print("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:
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!")
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:
# 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:
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);
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)
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)
approx # estimates of the root
f_of_x(approx, a, b, c) # function f(x) at estimates approach zero...
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).
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.