How to implement an efficient infinite generator of prime numbers in Python?

generator, primes, python

Solution

“If I have seen further…”

The `erat2` function from the cookbook can be further sped up (by about 20-25%):

erat2a

import itertools as it
def erat2a( ):
    D = {  }
    yield 2
    for q in it.islice(it.count(3), 0, None, 2):
        p = D.pop(q, None)
        if p is None:
            D[q*q] = q
            yield q
        else:
            # old code here:
            # x = p + q
            # while x in D or not (x&1):
            #     x += p
            # changed into:
            x = q + 2*p
            while x in D:
                x += 2*p
            D[x] = p

The `not (x&1)` check verifies that `x` is odd. However, since both `q` and `p` are odd, by adding `2*p` half of the steps are avoided along with the test for oddity.

erat3

If one doesn't mind a little extra fanciness, `erat2` can be sped up by 35-40% with the following changes (NB: needs Python 2.7+ or Python 3+ because of the `itertools.compress` function):

import itertools as it
def erat3( ):
    D = { 9: 3, 25: 5 }
    yield 2
    yield 3
    yield 5
    MASK= 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 1, 0, 0,
    MODULOS= frozenset( (1, 7, 11, 13, 17, 19, 23, 29) )

    for q in it.compress(
            it.islice(it.count(7), 0, None, 2),
            it.cycle(MASK)):
        p = D.pop(q, None)
        if p is None:
            D[q*q] = q
            yield q
        else:
            x = q + 2*p
            while x in D or (x%30) not in MODULOS:
                x += 2*p
            D[x] = p

The `erat3` function takes advantage of the fact that all primes (except for 2, 3, 5) modulo 30 result to only eight numbers: the ones included in the `MODULOS` frozenset. Thus, after yielding the initial three primes, we start from 7 and work only with the candidates. The candidate filtering uses the `itertools.compress` function; the “magic” is in the `MASK` sequence; `MASK` has 15 elements (there are 15 odd numbers in every 30 numbers, as chosen by the `itertools.islice` function) with a `1` for every possible candidate, starting from 7. The cycle repeats as specified by the `itertools.cycle` function. The introduction of the candidate filtering needs another modification: the `or (x%30) not in MODULOS` check. The `erat2` algorithm processed all odd numbers; now that the `erat3` algorithm processes only r30 candidates, we need to make sure that all `D.keys()` can only be such —false— candidates.

Benchmarks

Results

On an Atom 330 Ubuntu 9.10 server, versions 2.6.4 and 3.1.1+:

$ testit
up to 8192
==== python2 erat2 ====
100 loops, best of 3: 18.6 msec per loop
==== python2 erat2a ====
100 loops, best of 3: 14.5 msec per loop
==== python2 erat3 ====
Traceback (most recent call last):
…
AttributeError: 'module' object has no attribute 'compress'
==== python3 erat2 ====
100 loops, best of 3: 19.2 msec per loop
==== python3 erat2a ====
100 loops, best of 3: 14.1 msec per loop
==== python3 erat3 ====
100 loops, best of 3: 11.7 msec per loop

On an AMD Geode LX Gentoo home server, Python 2.6.5 and 3.1.2:

$ testit
up to 8192
==== python2 erat2 ====
10 loops, best of 3: 104 msec per loop
==== python2 erat2a ====
10 loops, best of 3: 81 msec per loop
==== python2 erat3 ====
Traceback (most recent call last):
…
AttributeError: 'module' object has no attribute 'compress'
==== python3 erat2 ====
10 loops, best of 3: 116 msec per loop
==== python3 erat2a ====
10 loops, best of 3: 82 msec per loop
==== python3 erat3 ====
10 loops, best of 3: 66 msec per loop

Benchmark code

A `primegen.py` module contains the `erat2`, `erat2a` and `erat3` functions. Here follows the testing script:

#!/bin/sh
max_num=${1:-8192}
echo up to $max_num
for python_version in python2 python3
do
    for function in erat2 erat2a erat3
    do
        echo "==== $python_version $function ===="
        $python_version -O -m timeit -c \
        -s  "import itertools as it, functools as ft, operator as op, primegen; cmp= ft.partial(op.ge, $max_num)" \
            "next(it.dropwhile(cmp, primegen.$function()))"
    done
done

Problem

This is not a homework, I am just curious. INFINITE is the key word here. I wish to use it as `for p in primes()`. I believe that this is a built-in function in Haskell. So, the answer cannot be as naive as "Just do a Sieve". First of all, you do not know how many consecutive primes will be consumed. Well, suppose you could concoct 100 of them at a time. Would you use the same Sieve approach as well as the frequency of prime numbers formula? I prefer non-concurrent approach. Thank you for reading (and writing ;) )!

Original source

Related problems