Segmented and wheel factorized prime generator in Python

Question

I recently discovered that it is much faster to generate primes using a Sieve of Eratosthenes (SoE) with a boolean array data structure than with a heap. Since then I have been looking for fast and scalable implementations. The fastest Python version I found is here in a comment by Willy Good.

But Willy's code is just to demonstrate how wheel factorization works, I believe. It uses O(n) memory. When I run it for n > 2.5e9 or so, my laptop with 8G of RAM starts to thrash (excessive paging).

I realize that using a segmented SoE makes it scalable, so I experimented with simple segmented sieves. That eliminated the thrashing for big N but was considerably slower than using mod 30 wheel factorization.

My next goal was to find a combination of wheel factorization and segmentation. Kim Walisch's primesieve is a great example in C++ with very helpful doc, and Gordon B Good has a fast javascript version, but I couldn't find anything for Python. Here is my version (sorry for the length):

#!/usr/bin/python3 -Wall
# program to find all primes up to and including n, using a segmented wheel sieve  

from sys import argv, stdout
from bitarray import bitarray

# Counts and optionally prints all prime numbers no larger than 'n'  

#CUTOFF      = 10          # for debugging only
#SIEVE_SIZE  = 2           # for debugging only
CUTOFF      = 1e4
SIEVE_SIZE  = 2**20
GHz         = 1.6          # on my i5-6285U laptop

# mod 30 wheel constant arrays
modPrms     = [7,11,13,17,19,23,29,31]
modPrmsM30  = [7,11,13,17,19,23,29,1]
gaps        = [4,2,4,2,4,6,2,6,4,2,4,2,4,6,2,6] # 2 loops for overflow
ndxs        = [0,0,0,0,1,1,2,2,2,2,3,3,4,4,4,4,5,5,5,5,5,5,6,6,7,7,7,7,7,7]
rnd2wh      = [7,7,0,0,0,0,0,0,1,1,1,1,2,2,3,3,3,3,4,4,5,5,5,5,6,6,6,6,6,6]

def num2ix(n):
    """Return the wheel index for n."""
    n = n - 7              # adjust for wheel starting at 1st prime past 2,3,5 vs. 0
    return (n//30 << 3) + ndxs[n % 30]       

def ix2num(i):
    """Return a number matching i (a wheel index)."""
    return 30 * (i >> 3) + modPrms[i & 7]   

def progress(j, num_loops, enabled):
    """Display a progress bar on the terminal."""
    if enabled:
        size = 60
        x = size*j//num_loops
        print("%s[%s%s] %i/%i\r" % ("Sieving: ", "#"*x, "."*(size-x), j, num_loops), end=' ')
        stdout.flush()

def prime_gen_wrapper(n):
    """Decide whether to use the segmented sieve or a simpler version.  Stops recursion."""
    if n < CUTOFF:
        return smallSieve(n+1) # rwh1 returns primes < N.  We need sieving primes <= sqrt(limit)
    else:
        return segmentedSieve(n)

def smallSieve(n):
    """Returns a list of primes less than n."""
    # a copy of Robert William Hanks' rwh1 used to get sieving primes for smaller ranges
    # https://stackoverflow.com/questions/2068372/fastest-way-to-list-all-primes-below-n-in-python/3035188#3035188
    sieve = [True] * (n//2)
    for i in range(3,int(n**0.5)+1,2):
        if sieve[i//2]:
            sieve[i*i//2::i] = [False] * ((n-i*i-1)//(2*i)+1)
    return [2] + [2*i+1 for i in range(1,n//2) if sieve[i]]

def segmentedSieve(limit, statsOnly = False): 
    """
    Sieves potential prime numbers up to and including limit.

    statsOnly (default False) controls the return.
        when False, returns a list of primes found.
        when True, returns a count of the primes found.
    """
    # segmentation originally based on Kim Walisch's simple C++ example of segmantation found here 
    #     https://github.com/kimwalisch/primesieve/wiki/Segmented-sieve-of-Eratosthenes

    # mod 30 wheel factorization based on a non-segmented version found here in a comment by Willy Good
    # https://programmingpraxis.com/2012/01/06/pritchards-wheel-sieve/  

    sqrt = int(limit ** 0.5)
    lmtbf = SIEVE_SIZE * 8
    while (lmtbf >> 1) >= limit:
       lmtbf >>= 1         # adjust the sieve size downward for small N

    multiples = []; wx = []
    outPrimes = [2,3,5]    # the wheel skips multiples of these, but they may be needed as output
    count = len(outPrimes) 
    lim_ix = num2ix(limit)
    buf = bitarray(lmtbf)
    show_progress = False
    if statsOnly:   # outer loop?
        print("sieve size:", end=' ')
        ss = len(memoryview(buf))
        if ss > 1024:
            print(ss//1024, "KB")
        else:
            print(ss, "bytes")
        if limit > 1e8:
            show_progress = True

    num_loops = (lim_ix + lmtbf - 1)//(lmtbf)   # round up

    # get sieving primes recursively, skipping those eliminated by the wheel
    svPrimes  = prime_gen_wrapper(sqrt)[count:] 

    for lo_ix in range(0, lim_ix + 1, lmtbf):   # loop over all the segments
        low = ix2num(lo_ix)
        high = ix2num(lo_ix + lmtbf) - 1
        buf.setall(True)
        progress(lo_ix//(lmtbf), num_loops, show_progress)

        # generate new multiples of sieving primes and wheel indices needed in this segment
        for p in svPrimes[len(multiples):]:
            pSquared = p * p
            if pSquared > high:
                break
            multiples.append(pSquared)
            wx.append(num2ix(p) & 7)

        # sieve the current segment
        for x in range(len(multiples)):
            s  = multiples[x]
            if s <= high:
                p  = svPrimes[x]
                ci = wx[x]
                s -= 7
                p8 = p << 3
                for j in range(8):
                    c = (s//30 << 3) + ndxs[s % 30] - lo_ix
                    # buf[c::p8] = False * ((lmtbf - c) // p8 + 1)
                    buf[c::p8] = False              # much simpler with bitarray vs. pure python
                    s += p * gaps[ci]; ci += 1

        # calculate the next multiple of p to sieve in an upcoming segment and its wheel index
                f       = (high + p - 1)//p         # next factor of a multiple of p past this segment
                f_mod   = f % 30
                i = rnd2wh[f_mod]                   # round up to next wheel index to eliminate multiples of 2,3,5
                nxt = p * (f - f_mod + modPrmsM30[i])   # back to a normal multiple of p past this segment
                wx[x] = i                               # save wheel index 
                multiples[x] = nxt                      #                  ... and next multiple of p

        # handle any extras in the last segment
        if high > limit:
            top = lim_ix - lo_ix
        else:
            top = lmtbf -1

        # collect results from this segment
        if statsOnly:
            count += buf[:top+1].count()
        else:
            for i in range(top + 1):
                if buf[i]:
                    x = i + lo_ix
                    p = 30 * (x >> 3) + modPrms[x & 7]   # ix2num(x) inlined, performance is sensitive here 
                    outPrimes.append(p)

    if show_progress:
        progress(num_loops, num_loops, True) 
        print()

    if statsOnly:
        return count
    else:
        return outPrimes

# Driver Code 
if len(argv) < 2:
    a = '1e8'
else:
    a = argv[1]

n = int(float(a))

from math import log
from time import time
#from datetime import timedelta

start = time()
count = segmentedSieve(n, statsOnly = True)

elapsed = time() - start
BigOculls = n * log(log(n,2),2)
cycles = GHz * 1e9 * elapsed 
cyclesPerCull = cycles/BigOculls

print(count, "primes found <=", a)
print("%.3f seconds, %.2f cycles per Big-O cull" %(elapsed, cyclesPerCull))

if count < 500:
    print(segmentedSieve(n))

Is anyone aware of another Python prime generator that is segmented and faster for large sizes? Any ideas to speed this one up, or make the code more compact or more clear? I had been using Willy Good's mod 30 unsegmented wheel sieve for smallSieve() here because it is faster, but Robert William Hank's primes_rwh1 is more compact and nearly as good for big N. I am not necessarily tied to using a mod 30 wheel; if someone is aware of a faster implementation and can demonstrate that it beats Willy's code with a benchmark, I am all ears.

If I didn't care somewhat about code size, I would implement some features found in Kim Walisch's primesieve, such as:

• pre_sieving for primes up to 19, then copying the result into each segment
• dividing up the sieving primes into small, medium, and large sizes, and processing each group differently

...but this is probably too long already.

Originally I wanted this to be pure Python but I realized that the bitarray package fit my needs well.

EDIT

Some benchmarks against Willy Good's unsegmented mod 30 wheel sieve, the fastest Python implementation I'm currently aware of for smaller sizes. Willy's is prime_wheel.py, the segmented wheel sieve is prime_ba.py (ba == bitarry, the last significant change). First at 1 million:

$ time ./prime_ba.py 1e6
sieve size:  1024 KB
78498 primes found <= 1e6
0.032 seconds, 11.68 cycles per Big-O cull

real    0m0.064s
user    0m0.031s
sys     0m0.000s
$ time ./prime_wheel.py 1e6
78498 primes found <= 1e6

real    0m0.053s
user    0m0.016s
sys     0m0.031s

The unsegmented wheel sieve is a little faster than my segmented version. But both run in under .1 sec so I'm not too worried. Next at 100 million:

$ time ./prime_ba.py 1e8
sieve size:  1024 KB
5761455 primes found <= 1e8
0.290 seconds, 0.98 cycles per Big-O cull

real    0m0.322s
user    0m0.297s
sys     0m0.016s
$ time ./prime_wheel.py 1e8
5761455 primes found <= 1e8

real    0m2.789s
user    0m2.500s
sys     0m0.281s

This is starting to show the effects of the different memory footprints. The segmented version is only using 1M of RAM for sieving, the unsegmented version uses O(n) memory. That is my incentive for creating this version. At 10 billion:

$ time ./prime_ba.py 1e10
sieve size:  1024 KB
Sieving: [############################################################] 318/318
455052511 primes found <= 1e10
33.420 seconds, 1.06 cycles per Big-O cull

real    0m33.451s
user    0m33.297s
sys     0m0.016s
$ time ./prime_wheel.py 1e10
^C^CTraceback (most recent call last):
  File "./prime_wheel.py", line 36, in <module>
    for x in primes235(n):
  File "./prime_wheel.py", line 22, in primes235
    buf[c::p8] = [False] * ((lmtbf - c) // p8 + 1)
KeyboardInterrupt

^C
real    3m16.165s
user    0m32.734s
sys     2m15.953s

The segmented version chugs along, still using a 1MB sieve. The unsegmented version uses all of my 8G of RAM, the system starts to page excessively, the fan shifts into high gear. I hit ctrl-C several times to get out of it after 3 minutes. The "sys" time is now dominant due to the paging.

EDIT 2

Replaced the code with a new version to:

• fix an off-by-one error calling smallSieve() a.k.a. rwh1_primes, which generates primes less than N. When it is used to generate sieving primes, we need to get all of the primes up to and including N, the integer square root of the input limit. External symptom: some composite numbers are reported as primes.
• shrink the bitarray when it is much larger than needed for the input limit. This results in a dramatic speedup for smaller sizes since the entire bitarray is always sieved to simplify the segmentation loop.
• report the sieve size in bytes when appropriate due to the previous change
• a few minor cleanups

If anyone is interested in seeing a diff of the changes, please let me know in the comments.

EDIT 3

• Replaced the code with a Python 3 version. "2to3-2.7" made the conversion much easier than I feared. Once 2to3 was done, about all I had to do was change "/" to "//" in bunches of places to get integer/floor division and test it. Thanks again @GZ0 for pointing out how soon Python 2.7 support is going away.
• Moved the code to calculate the total number of segments for the progress bar from the segmentation loop itself and into the the initialization (blush).
• Add some rudimentary docstrings.

EDIT 4

A new OO version incorporating changes suggested by @QuantumChris is available here.

Answer 1

Hi welcome to code review! Interesting topic, I remember writing some different prime sieves for project Euler problems.

Stylistically, it would really help to use PEP8. It's python's standard style-guide which helps readability for yourself and others. Things like imports at the top, consistent naming, etc.

There are a few places where I think ternary operators would really clean up the code: e.g.

if statsOnly:
    return count
else:
    return outPrimes

would be replaced with

return count if statsOnly else outPrimes

You have a lot of variables and code floating around outside of functions. I think a class would serve well to fix this. You could have your modPrms, modPrmsM30 etc as class or instance variables and the functions like num2ix() as methods of the class. A rough outline of the class might be:

class PrimeSieve:

    def __init__(self):
        self.cutoff = 1e4
        self.sieve_size = 2 ** 20
        self.clock_speed = 1.6  # In GHz

        # mod 30 wheel constant arrays
        self.mod_primes = [7, 11, 13, 17, 19, 23, 29, 31]
        self.mod_primes_m30 = [7, 11, 13, 17, 19, 23, 29, 1]
        self.gaps = [4, 2, 4, 2, 4, 6, 2, 6, 4, 2, 4, 2, 4, 6, 2, 6]  # 2 loops for overflow
        self.ndxs = [0, 0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 7, 7, 7, 7, 7, 7]
        self.rnd2wh = [7, 7, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6]

    def num2ix(self, n):
        """Return the wheel index for n."""
        # Adjust for wheel starting at 1st prime past 2,3,5 vs. 0
        n -= 7
        return (n // 30 << 3) + self.ndxs[n % 30]

    ...

You could also provide things like clock_speed as arguments which might be preferred (just put these into the init arguments):

def __init__(self, cutoff, sieve_size, clock_speed):
    ...

It seems weird to me to have progress() contain an enabled argument which basically decides whether anything is done at all. I would remove this argument, and simply wrap the function call with an if. For displaying progress I'd also highly recommend using tqdm which is made for exactly this kind of thing. print() also has a flush argument which will flush output. If you don't want to use tqdm, switch to using f-strings or .format() which are much more readable than the old % style you're using.

You can add file """docstrings""" just as you have function docstrings. These sit at the top of the file and are preferred over introductory comments.

Timing functions and methods is often done well using decorators. These wrap methods allowing you to execute code before and after their execution which is helpful for timing, logging and all sorts of other things. The following is a simple example I use a lot. It can be applied to functions and methods:

from functools import wraps


def timer_func(orig_func):
    """
    Prints the runtime of a function when applied as a decorator (@timer_func).
    """

    @wraps(orig_func)
    def wrapper(*args, **kwargs):

        t1 = time()
        result = orig_func(*args, **kwargs)
        t2 = time() - t1

        print(f'{orig_func.__qualname__} ran in: {t2} seconds.')

        return result
    return wrapper

You could write another decorator which counts the number of calls of a function, see here.

Your variable naming could be much improved. It should be obvious what everything is. GHz -> clock_speed; modPrms -> mod_primes rnd2wh -> literally anything else. Using i, j or x is fine for small one-off index names or iterables but not for such huge sections of code.

The variable low is declared but not used. This may be a bug.

If you want to iterate over an object and get its indices, use enumerate():

for i, multiple in enumerate(multiples):
    ...

segmented_sieve() should really be broken up. You have a lot of code here for processing, printing, formatting ... Try to have your functions perform single, short tasks. This also makes it much easier to convert functions to generators as you don't need to jump in and out, you can often just swap a return for a yield and call it as an iterable. Modularity also helps with readability, debugging, testing and extending.

It's recommended to wrap code you call in if __name__ == '__main__': See this excellent answer for why.

There's more to be said but I have to go for now; I may add more comments later. Feel free to post another question with the above changes where you may get more specific feedback on optimisations and such.