Iterators and Generators - Scientific Computing with Python

So far, you have seen things like

a = [1,2,3]
for x in a:
    print(x)

1
2
3

from itertools import permutations, product, repeat

import numpy as np

This looks very different from a C-style for loop where we loop over the variable index:

for (size_t i = 0; i < 3; i++) {
    printf("%d\n", i);
}

Or for instance, we can use something called a range

for i in range(3):
    print(i)

0
1
2

or other data types

d = {"hello": 1, "goodbye": 2}
for k,v in d.items():
    print(k, ' : ', v)

hello  :  1
goodbye  :  2

The key to using this sort of syntax is the concept of iterator. This is common in object-oriented programming (not just in Python), but you probably haven’t seen iterators before if you’ve only used imperative languages.

An object is iterable if it implements the __iter__ method, which is expected to return an iterator object. An object is an iterator if it implements the __next__ method, which either

returns the next element of the iterable object
raises the StopIteration exception if there are no more elements to iterate over

A Basic Iterator¶

What if we want to replicate range?

r = range(3)
type(r)

range

we can produce an iterator using the iter function

ri = iter(r)
type(ri)

range_iterator

we can explicitly run through the iterator using the next function

next(ri)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-11-e29b1d0ccf05> in <module>
----> 1 next(ri)

StopIteration:

r = range(1,5,2)
ri = iter(r)
print(next(ri))
print(next(ri))

1
3

class my_range_iterator(object):
    def __init__(self, start, stop, stride):
        self.state = start
        self.stop = stop
        self.stride = stride
        
    def __next__(self):
        if self.state >= self.stop:
            raise StopIteration  # signals "the end"
        ret = self.state # we'll return current state
        self.state += self.stride # increment state
        return ret
        
        
# an iterable
class my_range(object):
    def __init__(self, start, stop, stride=1):
        self.start = start
        self.stop = stop
        self.stride = stride
    
    def __iter__(self):
        return my_range_iterator(self.start, self.stop, self.stride)

r = my_range(0,3)
type(r)

__main__.my_range

ri = iter(r)
type(ri)

__main__.my_range_iterator

next(ri)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-20-e29b1d0ccf05> in <module>
----> 1 next(ri)

<ipython-input-14-f6ee4fb10d39> in __next__(self)
      7     def __next__(self):
      8         if self.state >= self.stop:
----> 9             raise StopIteration  # signals "the end"
     10         ret = self.state # we'll return current state
     11         self.state += self.stride # increment state

StopIteration:

for i in my_range(0,3):
    print(i)

0
1
2

for i in range(0,3):
    print(i)

0
1
2

You can also create classes that are both iterators and iterables

# an iterable and iterator
class my_range2(object):
    def __init__(self, start, stop, stride=1):
        self.start = start
        self.stop = stop
        self.stride = stride
        self.state = start
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.state >= self.stop:
            raise StopIteration  # signals "the end"
        ret = self.state # we'll return current state
        self.state += self.stride # increment state
        return ret

Using Iterators for Computation¶

Let’s now use iterators for something more interesting - computing the Fibonacci numbers.

class FibonacciIterator(object):
    def __init__(self):
        self.a = 0 # current number
        self.b = 1 # next number
        
    def __iter__(self):
        return self
    
    def __next__(self):
        ret = self.a
        self.a, self.b = self.b, self.a + self.b # advance the iteration
        return ret

for i in FibonacciIterator():
    if i > 1000:
        break
    print(i)

Note that we never raise a StopIteration exception - the iterator will just keep going if we let it.

Exercise¶

Define FibonacciIterator so it will iterate over all Fibonacci numbers until they are greater than a parameter n.

## Your code here

Generators¶

Often, a more elegant way to define an iterator is using a generator

This is a special kind of iterator defined using a function instead of using classes that implement the __iter__ and __next__ methods.

See this post for more discussion.

def my_range3(state, stop, stride=1):
    while state < stop:
        yield state
        state += stride

Note that we use the def keyword instead of the class keyword for our declaration. The yield keyword returns subsequent values of the iteration.

r = my_range3(0,3)
type(r)

generator

ri = iter(r)
type(ri)

generator

next(ri)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-42-e29b1d0ccf05> in <module>
----> 1 next(ri)

StopIteration:

for i in my_range3(0,3):
    print(i)

0
1
2

Our Fibonacci example re-written using a generator:

def FibonacciGenerator():
    a = 0
    b = 1
    while True:
        yield a
        a, b = b, a + b

for i in FibonacciGenerator():
    if i > 1000:
        break
    print(i)

Exercise¶

Define FibonacciGenerator so it will iterate over all Fibonacci numbers until they are greater than a parameter n.

## Your code here

def FibonacciGenerator(n):
    a = 0
    b = 1
    while a < n:
        yield a
        a, b = b, a + b
        
for i in FibonacciGenerator(1000):
    print(i)

Iteration tools¶

Some useful tools for iterators that come in handy are:

zip - iterates over multiple iterators simulataneously

for i, a in zip([0,1,2], ['a', 'b', 'c']):
    print(i,a)

0 a
1 b
2 c

reversed - iterates in reverse order

for i in reversed(range(3)):
    print(i)

2
1
0

enumerate - returns the iteration step count as well as the iterator value

for i, a in enumerate('abc'):
    print(i, a)

0 a
1 b
2 c

Exercise¶

Implement your own versions of zip and enumerate using generators

## Your code here

def my_zip(a, b):
    ai = iter(a)
    bi = iter(b)
    while True:
        try:
            yield next(ai), next(bi)
        except:
            return
        
for i, a in my_zip(range(3), 'abc'):
    print(i,a)

0 a
1 b
2 c

The Itertools Package¶

A useful package for dealing with iterators is the itertools package. Here are a few examples - click on the link to see what else the package provides.

product gives the equivalent of a nested for-loop

from itertools import product

for i, j in product(range(2), range(3)):
    print(i, j)

# equivalent to:
for i in range(2):
    for j in range(3):
        print(i,j)

repeat just repeats a value

from itertools import repeat

for i in repeat(10, 5):
    print(i)

permutations

from itertools import permutations

for p in permutations(range(3)):
    print(p)

(0, 1, 2)
(0, 2, 1)
(1, 0, 2)
(1, 2, 0)
(2, 0, 1)
(2, 1, 0)

Exercise¶

Implement your own version of product and repeat using generators.

## Your code here

Iterators for Scientific Computing¶

One way you might use an iterator in scientific computing is when implementing an iterative algorithm.

Here is an example of the power method, which finds the largest eigenvalue-eigenvector pair of a matrix.

import numpy as np

def PowerMethodGenerator(A, x):
    
    def RayleighQuotient(A, x):
        """
        x^T A x / x^T x
        """
        return np.dot(x, A @ x) / np.dot(x,x)
    
    x = x / np.linalg.norm(x)
    rq_prev = np.inf
    rq = RayleighQuotient(A, x)
    
    while True:
        # yield state: RayleighQuotient, x, and difference from previous iteration
        yield rq, x, np.abs(rq - rq_prev)
        
        # compute next iteration
        x = A @ x
        x = x / np.linalg.norm(x)
        rq_prev = rq
        rq = RayleighQuotient(A, x)

# here's a version that uses the generator in a while-loop

n = 100
A = np.random.randn(n, n)
A = A @ A.T + 5 # constant increases eigenvalue in constant vector direction
x0 = np.random.randn(n)

solver = PowerMethodGenerator(A, x0)
tol = 1e-4

while True:
    rq, x, eps = next(solver)
    print(rq, eps)
    if eps < tol:
        break

89.81253501091724 inf
242.85910417948966 153.04656916857243
444.85274716457826 201.9936429850886
574.2764337134453 129.423686548867
605.3333719186537 31.05693820520844
611.5407739473036 6.207402028649881
612.9106617975037 1.3698878502001435
613.2485808429876 0.3379190454838863
613.3392058160957 0.09062497310810613
613.3650279556526 0.02582213955690804
613.3727234642167 0.0076955085640975085
613.3750960037759 0.002372539559132747
613.375846586167 0.0007505823911060361
613.3760887640708 0.00024217790382863313
613.376168092833 7.932876224003849e-05

If we decide that we’re not satisfied with convergence yet, we can resume where we left off

tol = 1e-6

while True:
    rq, x, eps = next(solver)
    print(rq, eps)
    if eps < tol:
        break

613.3761943849217 2.6292088705304195e-05
613.3762031806232 8.795701432973146e-06
613.3762061456915 2.965068347293709e-06
613.3762071517388 1.006047227747331e-06
613.3762074950506 3.4331185361224925e-07

You can do the same thing with for-loops

# here's a version that uses the generator in a for-loop

n = 100
A = np.random.randn(n, n)
A = (A @ A.T) + 5 # constant increases eigenvalue in constant vector direction
x0 = np.random.randn(n)

solver = PowerMethodGenerator(A, x0)
tol = 1e-4

for rq, x, eps in solver:
    print(rq, eps)
    if eps < tol:
        break

tol = 1e-6
print('\nresuming iteration after decreasing tolerance\n')
for rq, x, eps in solver:
    print(rq, eps)
    if eps < tol:
        break

98.54240658228312 inf
263.89481933418114 165.35241275189802
323.65196672338845 59.75714738920732
390.0128085671451 66.36084184375665
484.9275155474049 94.91470698025978
563.0011106312606 78.07359508385576
599.1416496541569 36.14053902289629
611.5563409103167 12.414691256159813
615.457347021376 3.901006111059246
616.6819177841095 1.2245707627334923
617.078052353799 0.3961345696894796
617.2115459496376 0.13349359583867226
617.2585339675145 0.04698801787685625
617.275774701485 0.017240733970538713
617.2823363892841 0.006561687799035099
617.2849102204586 0.0025738311745726605
617.2859439103643 0.0010336899056255788
617.2863664706344 0.0004225602701808384
617.2865414495121 0.00017497887768058717
617.2866145761685 7.312665638892213e-05

resuming iteration after decreasing tolerance

617.2866453354011 3.075923257256363e-05
617.2866583321916 1.2996790474062436e-05
617.286663841021 5.508829417522065e-06
617.2866661810857 2.34006472510373e-06
617.2866671766113 9.95525624603033e-07