Skip to content

Instantly share code, notes, and snippets.

@MattOates

MattOates/collatz_numba.py

Last active February 26, 2020 13:41
Show Gist options
  • Save MattOates/8663519544fb4696c53e2005d2ecd062 to your computer and use it in GitHub Desktop.
Save MattOates/8663519544fb4696c53e2005d2ecd062 to your computer and use it in GitHub Desktop.
Download ZIP
Helper functions for investigating the collatz graph and paths
Raw
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)
@MattOates
Copy link
Author

MattOates commented Feb 26, 2020 •
edited
Loading 

  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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment