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

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


“If I have seen further…”

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

erat2a

import itertools as itdef 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 itdef 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+:

$ testitup 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:

$ testitup 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/shmax_num=${1:-8192}echo up to $max_numfor python_version in python2 python3do    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()))"    donedone


Since the OP asks for an efficient implementation, here's a significant improvement to the active state 2002 code by David Eppstein/Alex Martelli (seen here in his answer): don't record a prime's info in the dictionary until its square is seen among the candidates. Brings space complexity down to below O(sqrt(n)) instead of O(n), for n primes produced ( π(sqrt(n log n)) ~ 2 sqrt(n log n) / log(n log n) ~ 2 sqrt(n / log n) ). Consequently, time complexity is also improved, i.e. it runs faster.

Creates a "sliding sieve" as a dictionary of current multiples of each base prime (i.e. below the sqrt of the current production point), together with their step values:

from itertools import count                                         # ideone.com/aVndFMdef postponed_sieve():                   # postponed sieve, by Will Ness          yield 2; yield 3; yield 5; yield 7;  # original code David Eppstein,     sieve = {}                           #   Alex Martelli, ActiveState Recipe 2002    ps = postponed_sieve()               # a separate base Primes Supply:    p = next(ps) and next(ps)            # (3) a Prime to add to dict    q = p*p                              # (9) its sQuare     for c in count(9,2):                 # the Candidate        if c in sieve:               # c's a multiple of some base prime            s = sieve.pop(c)         #     i.e. a composite ; or        elif c < q:               yield c                 # a prime             continue                      else:   # (c==q):            # or the next base prime's square:            s=count(q+2*p,2*p)       #    (9+6, by 6 : 15,21,27,33,...)            p=next(ps)               #    (5)            q=p*p                    #    (25)        for m in s:                  # the next multiple             if m not in sieve:       # no duplicates                break        sieve[m] = s                 # original test entry: ideone.com/WFv4f

(the older, original code here was edited to incorporate changes as seen in the answer by Tim Peters, below). see also this for a related discussion.

Similar 2-3-5-7 wheel-based code runs ~ 2.15x faster (which is very close to the theoretical improvement of 3/2 * 5/4 * 7/6 = 2.1875).


For posterity, here's a rewrite of Will Ness's beautiful algorithm for Python 3. Some changes are needed (iterators no longer have .next() methods, but there's a new next() builtin function). Other changes are for fun (using the new yield from <iterable> replaces four yield statements in the original. More are for readability (I'm not a fan of overusing ;-) 1-letter variable names).

It's significantly faster than the original, but not for algorithmic reasons. The speedup is mostly due to removing the original's add() function, doing that inline instead.

def psieve():    import itertools    yield from (2, 3, 5, 7)    D = {}    ps = psieve()    next(ps)    p = next(ps)    assert p == 3    psq = p*p    for i in itertools.count(9, 2):        if i in D:      # composite            step = D.pop(i)        elif i < psq:   # prime            yield i            continue        else:           # composite, = p*p            assert i == psq            step = 2*p            p = next(ps)            psq = p*p        i += step        while i in D:            i += step        D[i] = step