Adding wheel factorization to an indefinite sieve

Question

I’m modifying an indefinite sieve of Eratosthenes from here so it uses wheel factorization to skip more composites than its current form of just checking all odds.

I

Answer 1

This is the version that I had come up with. It's not as clean as Ness' but it works. I'm posting it so there's another example on how to use wheel factorization in case anyone comes by. I've left in the ability to choose what wheel size to use but it's easy to nail down a more permanent one - just generate the size you want and paste that into the code.

from itertools import count

def wpsieve():
    """prime number generator
    call this function instead of roughing or turbo"""
    whlSize = 11
    initPrms, gaps, c = wheel_setup(whlSize)

    for p in initPrms:
        yield p

    primes = turbo(0, (gaps, c))

    for p, x in primes:
        yield p

def prod(seq, factor=1):
    "sequence -> product"
    for i in seq: factor *= i
    return factor

def wheelGaps(primes):
    """returns list of steps to each wheel gap
    that start from the last value in primes"""
    strtPt = primes.pop(-1)  # where the wheel starts
    whlCirm = prod(primes)  # wheel's circumference

    # spokes are every number that are divisible by primes (composites)
    gaps = []  # locate where the non-spokes are (gaps)
    for i in xrange(strtPt, strtPt + whlCirm + 1, 2):
        if not all(map(lambda x: i%x, primes)): continue  # spoke 
        else: gaps.append(i)  # non-spoke

    # find the steps needed to jump to each gap (beginning from the start of the wheel)
    steps = []  # last step returns to start of wheel
    for i, j in enumerate(gaps):
        if i == 0: continue
        steps.append(int(j - gaps[i-1]))
    return steps

def wheel_setup(num):
    "builds initial data for sieve"
    initPrms = roughing(num)  # initial primes from the "roughing" pump
    gaps = wheelGaps(initPrms[:])  # get the gaps
    c = initPrms.pop(-1)  # prime that starts the wheel

    return initPrms, gaps, c

def roughing(end):
    "finds primes by trial division (roughing pump)"
    primes = [2]
    for i in range(3, end + 1, 2):
        if all(map(lambda x: i%x, primes)):
            primes.append(i)
    return primes

def turbo(lvl=0, initData=None):
    """postponed prime generator with wheels (turbo pump)
    Refs:  http://stackoverflow.com/a/10733621
           http://stackoverflow.com/a/19391111"""

    gaps, c = initData

    yield (c, 0)

    compost = {}  # found composites to skip
    # store as current value: (base prime, wheel index)

    ps = turbo(lvl + 1, (gaps, c))

    p, x = next(ps)
    psq = p*p
    gapS = len(gaps) - 1

    ix = jx = kx = 0  # indices for cycling the wheel

    def cyc(x): return 0 if x > gapS else x  # wheel cycler

    while True:
        c += gaps[ix]  # add next step on c's wheel
        ix = cyc(ix + 1)  # and advance c's index

        bp, jx = compost.pop(c, (0,0))  # get base prime and its wheel index

        if not bp:

            if c < psq:  # prime
                yield c, ix  # emit index for above recursive level
                continue
            else:
                jx = kx  # swap indices as a new prime comes up
                bp = p
                p, kx = next(ps)
                psq = p*p

        d = c + bp * gaps[jx]  # calc new multiple
        jx = cyc(jx + 1)

        while d in compost:
            step = bp * gaps[jx]
            jx = cyc(jx + 1)
            d += step

        compost[d] = (bp, jx)

leaving in the option for the wheel size also lets you see how quickly larger wheels don't do much. Below is testing code for how long it takes to generate the wheel of selected size and how fast the sieve is with that wheel.

import time
def speed_test(num, whlSize):

    print('-'*50)

    t1 = time.time()
    initPrms, gaps, c = wheel_setup(whlSize)
    t2 = time.time()

    print('2-{} wheel'.format(initPrms[-1]))
    print('setup time: {} sec.'.format(round(t2 - t1, 5)))

    t3 = time.time()      
    prm = initPrms[:]
    primes = turbo(0, (gaps, c))
    for p, x in primes:
        prm.append(p)
        if len(prm) > num:
            break
    t4 = time.time()

    print('run time  : {} sec.'.format(len(prm), round(t4 - t3, 5)))
    print('prime sum : {}'.format(sum(prm)))

for w in [5, 7, 11, 13, 17, 19, 23, 29]:
    speed_test(1e7-1, w)

Here's how it ran on my computer using PyPy (Python 2.7 compatible) when set to generate ten million primes:

2- 3 wheel
setup time:  0.0 sec.
run time  : 18.349 sec.
prime sum : 870530414842019
--------------------------------------------------
2- 5 wheel
setup time:  0.001 sec.
run time  : 13.993 sec.
prime sum : 870530414842019
--------------------------------------------------
2- 7 wheel
setup time:  0.001 sec.
run time  :  7.821 sec.
prime sum : 870530414842019
--------------------------------------------------
2- 11 wheel
setup time:  0.03 sec.
run time  :  6.224 sec.
prime sum : 870530414842019
--------------------------------------------------
2- 13 wheel
setup time:  0.011 sec.
run time  :  5.624 sec.
prime sum : 870530414842019
--------------------------------------------------
2- 17 wheel
setup time:  0.047 sec.
run time  :  5.262 sec.
prime sum : 870530414842019
--------------------------------------------------
2- 19 wheel
setup time:  1.043 sec.
run time  :  5.119 sec.
prime sum : 870530414842019
--------------------------------------------------
2- 23 wheel
setup time: 22.685 sec.
run time  :  4.634 sec.
prime sum : 870530414842019

Larger wheels are possible, but you can see they become rather long to set up. There's also the law of diminishing returns as the wheels get larger - not much point to go past the 2-13 wheel as they don't really make it that much faster. I also ended up running into a memory error past the 2-23 wheel (which had some 36 million numbers in its gaps list).