Iteration

Introduction

In Functions, we talked about how important it is to reduce duplication in your code by creating functions instead of copying-and-pasting. Reducing code duplication has three main benefits:

  1. It’s easier to see the intent of your code, because your eyes are drawn to what’s different, not what stays the same.

  2. It’s easier to respond to changes in requirements. As your needs change, you only need to make changes in one place, rather than remembering to change every place that you copied-and-pasted the code.

  3. You’re likely to have fewer bugs because each line of code is used in more places.

One tool for reducing duplication is functions, which reduce duplication by identifying repeated patterns of code and extract them out into independent pieces that can be easily reused and updated. Another tool for reducing duplication is iteration, which helps you when you need to do the same thing to multiple inputs: repeating the same operation on different columns, or on different datasets.

In this chapter you’ll learn about iteration in three ways: explicit iteration, using for loops and while loops; iteration via comprehensions (eg list comprehensions); and iteration for polars data frames.

Prerequisites

This chapter will use the polars data analysis package.

For Loops

A loop is a way of executing a similar piece of code over and over in a similar way.

A for loop does something for the time that the condition is satisfied. For example,

name_list = ["Lovelace", "Smith", "Pigou", "Babbage"]

for name in name_list:
    print(name)
Lovelace
Smith
Pigou
Babbage

prints out a name until all names have been printed out.

Every for loop has three components:

  1. The output, here a print statement. But you can imagine a for loop that populates each entry of a data frame or list (but you should always create the full Python object first and populate it later rather than changing its size within the loop because the latter is slow).

  2. The sequence: for name in name_list:. This determines what to loop over: each run of the for loop will assign name to a different value from the iterable name_list. It doesn’t have to be a list, any iterable object will do. It’s useful to think of name above as a pronoun, like “it”.

  3. The body: print(name). This is the code that does the work. It’s run repeatedly, each time with a different value for name. The first iteration will effectively run print(name_list[0]), the second will run print(name_list[1]), and so on.

As long as your object is an iterable (ie you can iterate over it), then it can be used in this way in a for loop. The most common examples are lists and tuples, but you can also iterate over strings (in which case each character is selected in turn). One gotcha to be aware of is if you iterate over a string, say “hello”, instead of iterating over a list (or tuple) of strings, eg ["hello"]. In the latter case, you get:

for entry in ["hello"]:
    print(entry)
    print("---end entry---")
hello
---end entry---

While in the former you get something quite different and typically not all that useful:

for entry in "hello":
    print(entry)
    print("---end entry---")
h
---end entry---
e
---end entry---
l
---end entry---
l
---end entry---
o
---end entry---
TipExercise

Write a for loop that prints out “Python for Data Science” so that each word is printed in a successive iteration.

A useful trick with for loops is the enumerate() keyword, which runs through an index that keeps track of the place of items in a list:

name_list = ["Lovelace", "Smith", "Hopper", "Babbage"]

for i, name in enumerate(name_list):
    print(f"The name in position {i} is {name}")
The name in position 0 is Lovelace
The name in position 1 is Smith
The name in position 2 is Hopper
The name in position 3 is Babbage

Remember, Python indexes from 0 so the first entry of i will be zero. But, if you’d like to index from a different number, you can:

for i, name in enumerate(name_list, start=1):
    print(f"The name in position {i} is {name}")
The name in position 1 is Lovelace
The name in position 2 is Smith
The name in position 3 is Hopper
The name in position 4 is Babbage

Another useful pattern when doing for loops with dictionaries is iteration over key, value pairs. We’ll get to learn more about dictionaries very shortly, but for now what’s important is that they map a key to a value, for example “apple” might map to “fruit”. Let’s take our example from earlier that mapped cities to temperatures. If we wanted to iterate over both keys and values, we can write a for loop like this:

cities_to_temps = {"Paris": 28, "London": 22, "Seville": 36, "Wellesley": 29}

for key, value in cities_to_temps.items():
    print(f"In {key}, the temperature is {value} degrees C today.")
In Paris, the temperature is 28 degrees C today.
In London, the temperature is 22 degrees C today.
In Seville, the temperature is 36 degrees C today.
In Wellesley, the temperature is 29 degrees C today.

Note that we added .items() to the end of the dictionary. And note that we didn’t have to call the key key, or the value value: these are set by their position. But part of best practice in writing code is that there should be no surprises, and writing key, value makes it really clear that you’re using values from a dictionary.

TipExercise

Write a dictionary that maps four cities you know into their respective countries and print the results using the key, value iteration trick.

Another useful type of for loop is provided by the zip() function. You can think of the zip() function as being like a zipper, bringing elements from two different iterators together in turn. Here’s an example:

first_names = ["Ada", "Adam", "Grace", "Charles"]
last_names = ["Lovelace", "Smith", "Hopper", "Babbage"]

for forename, surname in zip(first_names, last_names):
    print(f"{forename} {surname}")
Ada Lovelace
Adam Smith
Grace Hopper
Charles Babbage

The zip function is super useful in practice.

TipExercise

Zip together the first names from above with this jumbled list of surnames: ['Babbage', 'Hopper', 'Smith', 'Lovelace'].

(Hint: you have seen a trick to help re-arrange lists earlier on in the Chapter.)

List (and Other) Comprehensions

There’s a second way to do loops in Python and, in most but not all cases, they run faster. More importantly, and this is the reason it’s good practice to use them where possible, they are very readable. They are called list comprehensions.

List comprehensions can combine what a for loop and (if needed) what a condition do in a single line of code. First, let’s look at a for loop that adds one to each value done as a list comprehension (NB: in practice, we would use super-fast numpy arrays for this kind of operation):

num_list = range(50, 60)
[1 + num for num in num_list]
[51, 52, 53, 54, 55, 56, 57, 58, 59, 60]

The general pattern is a bit similar to with the for loop but there are some differences. There’s no colon, and no indenting. The syntax is “do something with x” then for x in iterable. Finally, the expression is wrapped in a [ and ] to make the output a list.

Note that lists are not the only wrapping you can provide to this kind of structure. A ( and ) to make it a generator (don’t worry about what this is for now), a { and } to make it a set (an object that only contains unique values), or it’s possible to create a dictionary from a comprehension too! List comprehensions are the most common, so if you only remember one kind, remember them.

TipExercise

Create a list comprehension that multiplies numbers in the range from 1 to 10 by 5.

Did you get the range right?

Let’s now see how to include a condition within a list comprehension. Say we had a list of numbers and wanted to filter it according to whether the numbers divided by 3 or not using the modulo operator:

number_list = range(1, 40)
divide_list = [x for x in number_list if x % 3 == 0]
print(divide_list)
[3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39]

The syntax here is do something to x for x in something if x satisfies some condition.

Here’s another example that picks out only the names that include ‘Smith’ in them:

names_list = ["Joe Bloggs", "Adam Smith", "Sandra Noone", "leonara smith"]
smith_list = [x for x in names_list if "smith" in x.lower()]
print(smith_list)
['Adam Smith', 'leonara smith']

Note how we used ‘smith’ rather than ‘Smith’ and then used lower() to ensure we matched names regardless of the case they are written in.

We can even do a whole if … else construct inside a list comprehension:

names_list = ["Joe Bloggs", "Adam Smith", "Sandra Noone", "leonara smith"]
smith_list = [x if "smith" in x.lower() else "Not Smith!" for x in names_list]
print(smith_list)
['Not Smith!', 'Adam Smith', 'Not Smith!', 'leonara smith']

Many of the constructs we’ve seen can be combined. For instance, there is no reason why we can’t have a nested or repeated list comprehension using zip(), and, perhaps more surprisingly, sometimes these are useful!

first_names = ["Ada", "Adam", "Grace", "Charles"]
last_names = ["Lovelace", "Smith", "Hopper", "Babbage"]
names_list = [x + " " + y for x, y in zip(first_names, last_names)]
print(names_list)
['Ada Lovelace', 'Adam Smith', 'Grace Hopper', 'Charles Babbage']

An even more extreme use of list comprehensions can deliver nested structures:

first_names = ["Ada", "Adam"]
last_names = ["Lovelace", "Smith"]
names_list = [[x + " " + y for x in first_names] for y in last_names]
print(names_list)
[['Ada Lovelace', 'Adam Lovelace'], ['Ada Smith', 'Adam Smith']]

This gives a nested structure that (in this case) iterates over first_names first, and then last_names. (Note that this object is a list of lists of strings!)

Let’s see a dictionary comprehension now. These look a bit similar to set comprehensions because they use { and } at either end but they are different because they come with a colon separating the keys from the values:

{key: value for key, value in zip(first_names, last_names)}
{'Ada': 'Lovelace', 'Adam': 'Smith'}
TipExercise

Create a nested list comprehension that results in a list of lists of strings equal to [['a0', 'b0', 'c0'], ['a1', 'b1', 'c1'], ['a2', 'b2', 'c2']] (ie a combination of the first three integers and letters of the alphabet). You may find that you need to convert numbers to strings using str(x) to do this.

If you’d like to learn more about list comprehensions, check out these short video tutorials.

While Loops

while loops continue to execute code until their conditional expression evaluates to False. (Of course, if it evaluates to True forever, your code will just continue to execute…)

n = 10
while n > 0:
    print(n)
    n -= 1

print("execution complete")
10
9
8
7
6
5
4
3
2
1
execution complete

NB: in case you’re wondering what -= does, it’s a compound assignment that sets the left-hand side equal to the left-hand side minus the right-hand side.

You can use the keyword break to break out of a while loop, for example if it’s reached a certain number of iterations without converging.

TipExercise

Making use of import string and then string.ascii_lowercase to get the characters in the alphabet, write a while loop that iterates backwards through the alphabet (starting at “z”) before printing “execution complete”.

Iteration with polars Data Frames

For loops, while loops, and comprehensions can be used with data frames, but in Polars, they are even more strongly discouraged than in pandas. Polars is built on a columnar, vectorized, and expression-based engine, so row-by-row iteration breaks performance and prevents optimizations.

These built-in methods for iteration have an overlap with what we’ve seen in Data Transformation but we’ll dig a little deeper into assign()/assignment operations, apply(), and eval() here.

Instead of iterating, Polars encourages you to use expressions and lazy evaluation, which are much faster and more memory efficient.

Assignment Operations and assign

An assignment is a statement that assigns the value on the right to the object on the left with an equals sign in the middle.

Let’s imagine we have a data frame like this:

import numpy as np
import polars as pl

df = pl.DataFrame(np.random.normal(size=(6, 4)), schema=["a", "b", "c", "d"])
df
shape: (6, 4)
a b c d
f64 f64 f64 f64
1.820339 0.151898 0.68888 -0.551492
1.523806 0.293586 0.875301 0.185056
-1.868947 -0.555708 0.665352 0.636098
-0.578884 -1.298526 -1.188968 0.210889
0.764905 -0.036629 0.888739 0.952502
2.664783 0.781417 1.68281 -1.506762

polars has built-in expressions designed to operate over columns and rows. For example, to compute the median:

df.select(pl.all().median())
shape: (1, 4)
a b c d
f64 f64 f64 f64
1.144355 0.057635 0.782091 0.197972 
df.select(pl.concat_list(pl.all()).list.median().alias("row_median"))
shape: (6, 1)
row_median
f64
0.420389
0.584444
0.040195
-0.883926
0.826822
1.232114

In these cases, and others using built-in functions, the iteration is hidden. What if we want to do something that isn’t a built in and also isn’t an aggregation though? Let’s take the example of adding five to every entry. We could do it by explicitly iterating row by row, then repeat that for each column, ie

# Do not do this!


def add_five_slow(df):
    for i in range(len(df)):
        for j in range(len(df.columns)):
            df[i, j] = df[i, j] + 5


%timeit add_five_slow(df)
459 μs ± 7.72 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

But to do this, every individual cell must be accessed and operated on—so it is very slow, taking milliseconds. polars has far faster ways of performing the same operation. For simple operations on data frames with consistent type, you can simply add five to the whole data frame:

%timeit df + 5
52.7 μs ± 1.71 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

This took tens of microseconds, much faster.

This also works on a per column basis, so you can do df.with_columns(pl.col("a") + 5) and so on.

These operations have equivalents using method chaining; stringing multiple operations together. The version of df.with_columns(new_a = pl.col("a") + 5) would be:

df = df.with_columns(new_a=pl.col("a") + 5)

Expressions (Polars’ Alternative to apply)

What happens if you have a more complicated operation you want to perform? In pandas, you might reach for apply(). In polars, you almost never need an equivalent because its expression API is incredibly expressive.

Most “complicated” operations can be expressed directly using polars’ built-in expressions:

# Don't do this (slow, row-wise)
mean_new_a = df.select(pl.col("new_a").mean()).item()
df.with_columns(
    result=pl.struct(["a", "b", "c"]).map_elements(
        lambda x: x["a"] - mean_new_a * x["c"] / x["b"], return_dtype=pl.Float64
    )
)

# Do this instead (fast, vectorized)
df.with_columns(result=pl.col("a") - pl.col("new_a").mean() * pl.col("c") / pl.col("b"))
shape: (6, 6)
a b c d new_a result
f64 f64 f64 f64 f64 f64
40556.820339 40555.151898 40555.68888 40554.448508 40561.820339 -4.437717
40556.523806 40555.293586 40555.875301 40555.185056 40561.523806 -4.778987
40553.131053 40554.444292 40555.665352 40555.636098 40558.131053 -8.811196
40554.421116 40553.701474 40553.811032 40555.210889 40559.421116 -6.409462
40555.764905 40554.963371 40555.888739 40555.952502 40560.764905 -5.881595
40557.664783 40555.781417 40556.68281 40553.493238 40562.664783 -3.95772

The first expression would work, but it evaluates the computation row by row using a python lambda, which is slow and prevents polars from optimizing the query. The second approach uses native expressions, allowing polars to execute the computation efficiently in a fully vectorized and optimized manner.

In polars, there’s no eval() — you use expressions directly instead:

df = df.with_columns((pl.col("a") / pl.col("new_a")).alias("ratio"))
df
shape: (6, 6)
a b c d new_a ratio
f64 f64 f64 f64 f64 f64
40556.820339 40555.151898 40555.68888 40554.448508 40561.820339 0.999877
40556.523806 40555.293586 40555.875301 40555.185056 40561.523806 0.999877
40553.131053 40554.444292 40555.665352 40555.636098 40558.131053 0.999877
40554.421116 40553.701474 40553.811032 40555.210889 40559.421116 0.999877
40555.764905 40554.963371 40555.888739 40555.952502 40560.764905 0.999877
40557.664783 40555.781417 40556.68281 40553.493238 40562.664783 0.999877

You can also create boolean columns the same way:

df = df.with_columns((pl.col("a") > 0.5).alias("a_gt_0.5"))
df
shape: (6, 7)
a b c d new_a ratio a_gt_0.5
f64 f64 f64 f64 f64 f64 bool
40556.820339 40555.151898 40555.68888 40554.448508 40561.820339 0.999877 true
40556.523806 40555.293586 40555.875301 40555.185056 40561.523806 0.999877 true
40553.131053 40554.444292 40555.665352 40555.636098 40558.131053 0.999877 true
40554.421116 40553.701474 40553.811032 40555.210889 40559.421116 0.999877 true
40555.764905 40554.963371 40555.888739 40555.952502 40560.764905 0.999877 true
40557.664783 40555.781417 40556.68281 40553.493238 40562.664783 0.999877 true