Ulam spiral sketcher using NumPy and tkinter optimisations

ulam.py

from itertools import count, cycle
from math import ceil, sqrt
from numpy import arange, array, full, ones, insert, where
from numpy import append as ap
from time import perf_counter as pf
from time import sleep
from tkinter import Canvas, Tk, constants

root = Tk()
# root.attributes('-fullscreen', True) # make fullscreen window
n, w, h = 1, 1300, 1300 # n=2 because 1 messes with the algo, starting number
canvas = Canvas(root, bg="black", width=w, height=h)
ct = canvas.create_text

text_size, text_font = 2, "Hack" # text size
pixels_to_jump = text_size/0.9 # not sure why this value works!
upto = 50_000 # 100M will take ~2.3GB RAM and 20s on powerful machine
ANIMATION_DELAY = 0.0125 # def: 0.05
PRINT_CHAR = '❖' # '×'

primes = arange(2,upto+1)
isprime = ones((upto-1),dtype=bool)
R, U, L, D = ((pixels_to_jump, 0), 0), ((0, -pixels_to_jump), 90), \
        ((-pixels_to_jump, 0), 180), ((0, pixels_to_jump), 270)

run_times = {}

# print_me = arange(1,upto)
# ^ uncomment for numbers, also switch the ct() in draw loop
print_me = full((len(primes)) + 1, ' ')

def prime6(): # new prime generation algo
    """Prime number generation algo"""

    time_it(1)
    for factor in primes[:int(sqrt(upto))]:
        if isprime[(factor)//2]: # if it is an even number
            isprime[((factor*3)-2)::factor]=0 # find a factor amongst the existing primes
    insert(primes[isprime],0,1) # no factors found, add to list of primes
    time_it(1)

def direction():
    """Legacy function, will be removed next version"""
    for _ in range(1, int(sqrt(upto))):
        # return cycle((R, U, L, D)) # movement sequence
        return cycle((L, D, R, U)) # movement sequence
# ((-pixels_to_jump, 0), 180), ((0, pixels_to_jump), 270)...

def distance():
    """Legacy function, will be removed next version"""
    # loop increments x every 2 loops
    for x in range(1, int(sqrt(upto))+1):
        for _ in (0, 1):
            yield x
# 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6...

def transform_array():
    """Star of the show. Matches, folds and trims"""
    global print_me, primes, isprime

    time_it(2)
    # list comprehension to match PRINT_CHAR if prime, otherwise ' '
    time_it(3)
    print_me = [PRINT_CHAR if isprime[_] else ' ' for _ in range(len(primes))]
    time_it(3)

    # Pre-made list of what used to be distance(), named so for readability later, sorry!
    # should be 'distance_list'
    d_l = [x for x in range(1, int(sqrt(upto))+1) for y in (0, 1)]

    time_it(5)
    # fold matrix into the required shape ([[1],[2],[3, 4],[5, 6]...])
    # although print_me[X:Y] looks like a mess, the maths is simple:
    # X = sum(d_l[0:i] = WHAT HAS BEEN READ ALREADY
    # Y = sum(d_l[0:i]) + d_l[i] = HOW FAR NEEDS TO BE READ FROM X 
    # putting these two together, we can cobble together the required output
    temp_array = [print_me[(sum(d_l[0:i])) : sum(d_l[0:i]) + d_l[i]] \
                            for i in range(0, len(d_l))]
    time_it(5)

    # trim empty elements from array
    time_it(6)
    print_me = [_ for _ in temp_array if len(_) != 0]
    time_it(6)

    temp_array = []
    time_it(2)

def time_it(index):
    """Timing function, takes int for index,
       retrieve values from array run_times[n]"""
    if index in run_times: # number exists, stop timer
        run_times[index] = (pf() - run_times[index]) * 1000
        return run_times[index] # return milliseconds

    else: # index does not exist
        if (0 > index) or (index > 10): # value error
            return "Enter value 1-10."
        else: # start timer
            run_times[index]=pf()

def draw_spiral():
    """Draw the spiral on screen"""
    text_x = w // 2; text_y = h // 2;

    time_it(7)
    for _, c, coords in zip(range(len(print_me)), distance(), direction()):
        delta_x, delta_y = coords[0][0], coords[0][1] 
        offset_x = ((len(print_me[_]) * delta_x) / 0.75)
        offset_y = ((len(print_me[_]) * delta_y) / 0.75)

        text_x += delta_x + offset_x
        text_y += delta_y + offset_y
        # this expression works for int array, DO NOT DELETE
        # ct(text_x, text_y, fill="white", text=(str(print_me[_])[1:-1]),\
        #                         font=(text_font, text_size), angle = coords[1])
        # this expression works for str arrays, DO NOT DELETE
        ct(text_x, text_y, fill="white", text=(' '.join(print_me[_])),\
            font=(text_font, text_size), angle = coords[1], justify=constants.CENTER)
        
        text_x += delta_x + offset_x
        text_y += delta_y + offset_y

        root.update_idletasks() # used to update canvas in realtime
        canvas.pack()
        sleep(ANIMATION_DELAY) # recom: (upto / (upto * 20)) i.e 0.05
    time_it(7)

    # canvas.pack()
    # canvas.update()
    # canvas.postscript(file="ulam.ps", colormode='color')
    root.mainloop()

def print_stats():
    """Print time stats on console"""
    ANIMATION_TIME = (len(print_me) * (ANIMATION_DELAY * 1000))

    print(f"\nTimes for {len(primes)+1:,} numbers;\n\t  {(len(primes[isprime])):,} primes:")
    # str(primes[isprime][-1:])[1:-1] is a ridiculous convolution,
    # to print the value without brackets, wrapped in int() to comma separate
    print(f"[1] primes  : {(run_times[1] * 1000):.2f}µs (max: {int(str(primes[isprime][-1:])[1:-1]):,})")
    print(f"[2] xform   : {run_times[2]:.2f}ms")
    print(f"[3]  --match: {run_times[3]:.2f}ms")
    print(f"[4]  --list : n/a") # {run_times[4] * 1000:.2f}µs")
    print(f"[5]  --fold : {run_times[5]:.2f}ms")
    print(f"[6]  --trim : {run_times[6]:.2f}ms")
    print(f"[7]  draw   : {(run_times[7] - ANIMATION_TIME):,.2f}ms (+{(run_times[7]/1000):.2f}s animation)")
    print(f"    TOTAL   : {(run_times[1] + run_times[2] + run_times[7]) - ANIMATION_TIME:,.2f}ms\n")

def main():
    prime6()
    transform_array()
    draw_spiral()

    ## Stats:
    # if foo:
    print_stats()

if __name__ == "__main__":
    main()

########################
##### DOWN & DIRTY #####
########################
## Records (50K, delay=0)
## v1.0-
# primes: 295.13µs
#  xform: 11840.98ms    80%
#   --match: 4702.20ms  --40%
#   -- fold: 6112.60ms  --52%
#   -- trim: 624.96ms
#   draw: 2.87s         20%
#  total: 14.71s
# ------------------------------
## v1.1- 207x faster
# [1] primes  : 270.85µs
# [2] xform   : 4.59ms      2574x
# [3]  --match: 2.91ms      1621x
# [4]  -- list: 14.25µs
# [5]  -- fold: 1.56ms      3918x
# [6]  -- trim: 0.09ms      6933x
# [7]  draw   : 66.30ms     43x (not even optimised yet)
#        TOTAL: 71.16ms     207x
# ------------------------------
## v1.1- a real challenge (upto=100M, delay=0)
# [1] primes  : 1,267,025.34µs
# [2] xform   : 8,951.02ms
# [3]  --match: 5,496.00ms
# [4]  -- list: 514.88µs
# [5]  -- fold: 3,298.45ms
# [6]  -- trim: 155.93ms
# [7]  draw   : 8,631.62ms
#     TOTAL   : 18,849.67ms
# -------------------------------

ulam_final.mp4