Helper functions for investigating the collatz graph and paths

collatz_numba.py

#!/usr/bin/env python3
import numpy as np
from numba import njit, uint64
from numba.typed import Dict

collatz_graph = Dict.empty(key_type=uint64, value_type=uint64)
collatz_path_lengths = Dict.empty(key_type=uint64, value_type=uint64)

def collatz(n: uint64, collatz_graph: Dict, collatz_path_lengths: Dict):
    # collatz_path_lengths[n] = collatz_iter(n, collatz_graph)
    collatz_iter(n, collatz_graph)
    return collatz_path(n, collatz_graph)

@njit
def collatz_path(n: uint64, collatz_graph: Dict):
    path_length: uint64 = collatz_path_lengths[n]
    path = np.zeros(path_length, dtype=uint64)
    n_i: uint64 = n
    path[0] = n
    for i in range(1, path_length):
        path[i] = collatz_graph[n_i]
        n_i = collatz_graph[n_i]
    return path

@njit
def collatz_iter(n: uint64, collatz_graph: Dict) -> uint64:
    k: uint64 = n
    path_length: unit64 = uint64(1)
    while (k != uint64(1)):
        k_i: uint64 = uint64(collatz_graph[k])
        if k_i == uint64(0):
            if k % uint64(2):
                k_i = (uint64(3)*k) + uint64(1)
            else:
                k_i = uint64(k / uint64(2))
            collatz_graph[k] = k_i
        path_length = path_length+uint64(1)
        k = k_i
    return path_length



print(collatz(uint64(10), collatz_graph, collatz_path_lengths))
#for n in range(10000000):
#    collatz(n)

  File "/Users/matt/test/collatz_numba.py", line 43, in <module>
    print(collatz(uint64(10), collatz_graph, collatz_path_lengths))
  File "/Users/matt/test/collatz_numba.py", line 11, in collatz
    collatz_iter(n, collatz_graph)
  File "/Users/matt/.pyenv/versions/3.6.9/lib/python3.6/site-packages/numba/dictobject.py", line 738, in impl
    raise KeyError()
KeyError
python ~/test/collatz_numba.py  1.10s user 0.20s system 84% cpu 1.532 total