fast-abelian-sandpile/sandpile_numba.py

45 lines
1.3 KiB
Python

"""Just compiling with numba."""
import numba
import numpy as np
@numba.njit(numba.void(numba.int64[:,:]))
def apply_gravity(terrain):
"""
$ python -m pyperf timeit --fast -s 'from examples.sandpile_numpy import main' 'main(10_000, False)'
Mean +- std dev: 19.2 ms +- 0.3 ms
$ python -m pyperf timeit --fast -s 'from examples.sandpile_numpy import main' 'main(100_000, False)'
Mean +- std dev: 1.83 sec +- 0.04 sec
2**16 grains : 1.6s
2**17 grains : 3.9s
2**18 grains : 13s
2**19 grains : 51s
2**20 grains : 3m5
"""
shape = np.shape(terrain)
while True:
did_someting = False
for x, y in np.ndindex(shape):
if terrain[x, y] >= 4:
div, terrain[x, y] = divmod(terrain[x][y], 4)
terrain[x - 1][y] += div
terrain[x + 1][y] += div
terrain[x][y + 1] += div
terrain[x][y - 1] += div
did_someting = True
if not did_someting:
return
def main(height, show=True):
width = int(height**0.5) + 1
terrain = np.zeros((width, width), dtype=np.int64)
terrain[width // 2, width // 2] = height
apply_gravity(terrain)
if show:
import sys
np.set_printoptions(threshold=sys.maxsize, linewidth=999)
print(terrain)