Mandelbrot numba examples

.gitignore

@@ -105,3 +105,4 @@ ENV/
 
 # Vim swap files
 .*.swp
+.DS_Store

examples/maths/mandelbrot_set_numpy/bench.yaml

@@ -0,0 +1,21 @@
+name: Mandelbrot Set (NumPy)
+description: |
+    Compute the Mandelbrot Set, a set of points on the complex plane which always remain bounded by a threshold value while solving the quadratic recurrence equation. It is an iterative problem which is also compute intensive and visual in nature.  
+    In this benchmark we learn how numba can accelerate the NumPy codes.  
+    Author - <a href="https://github.com/animator">Ankit Mahato</a>
+
+input_generator: numpy_array.py:input_generator
+xlabel: Max iteration
+validator: numpy_array.py:validator
+implementations:
+    - name: numpy_array
+      description: Implementation where a 3-dimensional numpy array is used to store the computed RBG pixel values.
+      function: numpy_array.py:mandelbrot
+    - name: numba_numpy_array
+      description: numba.njit decorator is applied on the NumPy implementation.
+      function: numba_numpy_array.py:mandelbrot
+    - name: numba_prange_numpy_array
+      description: numba.njit decorator is applied on the NumPy implementation. Numpy is primarily array designed to be as fast as possible on a single core, whereas numba automatically compiles a version which can run in parallel utilizing multiple threads if it contains reduction functions, array math functions and many more functions or assignments or operations. Explicit for loop parallelization is also performed using numba.prange().
+      function: numba_prange_numpy_array.py:mandelbrot
+
+baseline: numpy_array

examples/maths/mandelbrot_set_numpy/numba_numpy_array.py

@@ -0,0 +1,22 @@
+import math
+import numba as nb
+
+@nb.njit
+def mandelbrot(iters, pixels, bbox, width, max_iter):
+    for y in range(width):
+        for x in range(width):
+            c0 = complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                         bbox[1] + (bbox[3]-bbox[1])*y/width) 
+            c = 0
+            i = 1
+            while i < max_iter:
+                if abs(c) > 2:  
+                    log_iter = math.log(i) 
+                    pixels[y, x, 0] = 255*(1+math.cos(3.32*log_iter))//2
+                    pixels[y, x, 1] = 255*(1+math.cos(0.774*log_iter))//2
+                    pixels[y, x, 2] = 255*(1+math.cos(0.412*log_iter))//2 
+                    break
+                c = c * c + c0
+                i += 1
+            iters[y, x] = i                
+    return iters, pixels

examples/maths/mandelbrot_set_numpy/numba_prange_numpy_array.py

@@ -0,0 +1,22 @@
+import math
+import numba as nb
+
+@nb.njit(parallel = True)
+def mandelbrot(iters, pixels, bbox, width, max_iter):
+    for y in nb.prange(width):
+        for x in range(width):
+            c0 = complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                         bbox[1] + (bbox[3]-bbox[1])*y/width) 
+            c = 0
+            i = 1
+            while i < max_iter:
+                if abs(c) > 2:  
+                    log_iter = math.log(i) 
+                    pixels[y, x, 0] = 255*(1+math.cos(3.32*log_iter))//2
+                    pixels[y, x, 1] = 255*(1+math.cos(0.774*log_iter))//2
+                    pixels[y, x, 2] = 255*(1+math.cos(0.412*log_iter))//2
+                    break
+                c = c * c + c0
+                i += 1
+            iters[y, x] = i                
+    return iters, pixels

examples/maths/mandelbrot_set_numpy/numpy_array.py

@@ -0,0 +1,46 @@
+import numpy as np
+
+#### BEGIN: numpy
+import math
+
+def mandelbrot(iters, pixels, bbox, width, max_iter):
+    for y in range(width):
+        for x in range(width):
+            c0 = complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                         bbox[1] + (bbox[3]-bbox[1])*y/width) 
+            c = 0
+            i = 1
+            while i < max_iter:
+                if abs(c) > 2:  
+                    log_iter = math.log(i) 
+                    pixels[y, x, 0] = 255*(1+math.cos(3.32*log_iter))//2
+                    pixels[y, x, 1] = 255*(1+math.cos(0.774*log_iter))//2
+                    pixels[y, x, 2] = 255*(1+math.cos(0.412*log_iter))//2                    
+                    break
+                c = c * c + c0
+                i += 1
+            iters[y, x] = i                
+    return iters, pixels
+## END: numpy
+
+WIDTH = 600
+PLANE = (-2.0, -1.5, 1.0, 1.5)
+
+def validator(input_args, input_kwargs, impl_output):
+    actual_iters, actual_pixels  = impl_output
+    expected_iters, expected_pixels = mandelbrot(*input_args, **input_kwargs)
+    np.testing.assert_allclose(expected_iters, actual_iters)
+    np.testing.assert_allclose(expected_pixels, actual_pixels)
+
+def make_arrays(width):
+    iters = np.zeros((width, width), dtype=np.uint16)
+    pixels = np.zeros((width, width, 3), dtype=np.uint8)
+    return iters, pixels
+
+def input_generator():
+    for maxiter in [100, 200, 500, 1000]:
+        iters, pixels = make_arrays(WIDTH)
+        yield dict(category=("",),
+                x=maxiter,
+                input_args=(iters, pixels, PLANE, WIDTH, maxiter),
+                input_kwargs={})

examples/maths/mandelbrot_set_py/bench.yaml

@@ -0,0 +1,21 @@
+name: Mandelbrot Set (Pure Python)
+description: |
+    Compute the Mandelbrot Set, a set of points on the complex plane which always remain bounded by a threshold value while solving the quadratic recurrence equation. It is an iterative problem which is also compute intensive and visual in nature.  
+    In this benchmark we learn how numba can accelerate the pure python code for mandelbrot set computation.  
+    Author - <a href="https://github.com/animator">Ankit Mahato</a>
+
+input_generator: python_list.py:input_generator
+xlabel: Max iteration
+validator: python_list.py:validator
+implementations:
+    - name: python_list
+      description: Pure Python implementation where a list of list of RBG tuples is used to store the computed RBG pixel values.
+      function: python_list.py:mandelbrot
+    - name: numba_python_list
+      description: numba.njit decorator is applied on the pure Python implementation.
+      function: numba_python_list.py:mandelbrot
+    - name: numba_prange_python_list
+      description: numba.njit decorator is applied on the pure Python implementation along with explicit for loop parallelization using numba.prange().
+      function: numba_prange_python_list.py:mandelbrot
+
+baseline: python_list

examples/maths/mandelbrot_set_py/numba_prange_python_list.py

@@ -0,0 +1,20 @@
+import math
+import numba as nb
+
+@nb.njit(parallel = True)
+def mandelbrot(bbox, width, max_iter):     
+    pixels = [[(0, 0, 0) for j in range(width)] for i in range(width)]
+    for y in nb.prange(width):
+        for x in range(width):
+            c0 = complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                         bbox[1] + (bbox[3]-bbox[1])*y/width) 
+            c = 0
+            for i in range(1, max_iter): 
+                if abs(c) > 2: 
+                    log_iter = math.log(i) 
+                    pixels[y][x] = (int(255*(1+math.cos(3.32*log_iter))/2), 
+                                    int(255*(1+math.cos(0.774*log_iter))/2), 
+                                    int(255*(1+math.cos(0.412*log_iter))/2))
+                    break
+                c = c * c + c0
+    return pixels

examples/maths/mandelbrot_set_py/numba_python_list.py

@@ -0,0 +1,20 @@
+import math
+import numba as nb
+
+@nb.njit
+def mandelbrot(bbox, width, max_iter):     
+    pixels = [[(0, 0, 0) for j in range(width)] for i in range(width)]
+    for y in range(width):
+        for x in range(width):
+            c0 = complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                         bbox[1] + (bbox[3]-bbox[1])*y/width) 
+            c = 0
+            for i in range(1, max_iter): 
+                if abs(c) > 2: 
+                    log_iter = math.log(i) 
+                    pixels[y][x] = (int(255*(1+math.cos(3.32*log_iter))/2), 
+                                    int(255*(1+math.cos(0.774*log_iter))/2), 
+                                    int(255*(1+math.cos(0.412*log_iter))/2))
+                    break
+                c = c * c + c0
+    return pixels

examples/maths/mandelbrot_set_py/python_list.py

@@ -0,0 +1,36 @@
+import math
+
+#### BEGIN: python list
+def mandelbrot(bbox, width, max_iter):     
+    pixels = [[(0, 0, 0) for j in range(width)] for i in range(width)]
+    for y in range(width):
+        for x in range(width):
+            c0 = complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                         bbox[1] + (bbox[3]-bbox[1])*y/width) 
+            c = 0
+            for i in range(1, max_iter): 
+                if abs(c) > 2: 
+                    log_iter = math.log(i) 
+                    pixels[y][x] = (int(255*(1+math.cos(3.32*log_iter))/2), 
+                                    int(255*(1+math.cos(0.774*log_iter))/2), 
+                                    int(255*(1+math.cos(0.412*log_iter))/2))
+                    break
+                c = c * c + c0
+    return pixels 
+#### END: python list
+
+
+def validator(input_args, input_kwargs, impl_output):
+    actual_output = impl_output
+    expected_output = mandelbrot(*input_args, **input_kwargs)
+    assert actual_output == expected_output
+
+def input_generator():
+    for maxiter in [100, 200, 500, 1000]:
+        width = 600
+        plane = (-2.0, -1.5, 1.0, 1.5)
+        yield dict(category=("",),
+                   x=maxiter,
+                   input_args=(plane, width, maxiter),
+                   input_kwargs={})
+

examples/maths/mandelbrot_set_stencil/bench.yaml

@@ -0,0 +1,21 @@
+name: Mandelbrot Set (Stencil Kernel)
+description: |
+    Compute the Mandelbrot Set, a set of points on the complex plane which always remain bounded by a threshold value while solving the quadratic recurrence equation. It is an iterative problem which is also compute intensive and visual in nature.  
+    In this benchmark we learn more about Stencil kernel (both serial & parallel) implementations.  
+    Author - <a href="https://github.com/animator">Ankit Mahato</a>
+
+input_generator: numpy_array.py:input_generator
+xlabel: Max iteration
+validator: numpy_array.py:validator
+implementations:
+    - name: numpy_array
+      description: Implementation where a 3-dimensional numpy array is used to store the computed RBG pixel values.
+      function: numpy_array.py:mandelbrot
+    - name: numba_stencil
+      description: Stencil kernel implementation.
+      function: numba_stencil.py:mandelbrot
+    - name: numba_parallel_stencil
+      description: Stencil kernel implementation with parallel execution.
+      function: numba_parallel_stencil.py:mandelbrot
+
+baseline: numpy_array

examples/maths/mandelbrot_set_stencil/numba_parallel_stencil.py

@@ -0,0 +1,27 @@
+import numba as nb
+import numpy as np
+
+@nb.stencil 
+def iterskernel(a, max_iter):
+    c = 0
+    i = 1
+    while i < max_iter:
+        if abs(c) > 2:  
+            break
+        c = c * c + a[0, 0]
+        i += 1
+    return i
+
+@nb.njit(parallel= True)
+def mandelbrot(iters, pixels, bbox, width, max_iter):
+    arr = [[complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                      bbox[1] + (bbox[3]-bbox[1])*y/width) for x in range(width)] 
+              for y in range(width)]
+    c_arr = np.array(arr)
+    iters = iterskernel(c_arr, max_iter)
+    pix = iters.reshape((width, width, 1))*np.array([1,1,1])
+    pix = np.where(pix == max_iter, 
+                   0, 
+                   255*(1+np.cos(np.log(pix)*np.array([3.32, 0.774, 0.412])))//2)
+    pixels = pix.astype(np.uint8)
+    return iters, pixels

examples/maths/mandelbrot_set_stencil/numba_stencil.py

@@ -0,0 +1,27 @@
+import numba as nb
+import numpy as np
+
+@nb.stencil 
+def iterskernel(a, max_iter):
+    c = 0
+    i = 1
+    while i < max_iter:
+        if abs(c) > 2:  
+            break
+        c = c * c + a[0, 0]
+        i += 1
+    return i
+
+@nb.njit
+def mandelbrot(iters, pixels, bbox, width, max_iter):
+    arr = [[complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                      bbox[1] + (bbox[3]-bbox[1])*y/width) for x in range(width)] 
+              for y in range(width)]
+    c_arr = np.array(arr)
+    iters = iterskernel(c_arr, max_iter)
+    pix = iters.reshape((width, width, 1))*np.array([1,1,1])
+    pix = np.where(pix == max_iter, 
+                   0, 
+                   255*(1+np.cos(np.log(pix)*np.array([3.32, 0.774, 0.412])))//2)
+    pixels = pix.astype(np.uint8)
+    return iters, pixels

examples/maths/mandelbrot_set_stencil/numpy_array.py

@@ -0,0 +1,46 @@
+import numpy as np
+
+#### BEGIN: numpy
+import math
+
+def mandelbrot(iters, pixels, bbox, width, max_iter):
+    for y in range(width):
+        for x in range(width):
+            c0 = complex(bbox[0] + (bbox[2]-bbox[0])*x/width, 
+                         bbox[1] + (bbox[3]-bbox[1])*y/width) 
+            c = 0
+            i = 1
+            while i < max_iter:
+                if abs(c) > 2:  
+                    log_iter = math.log(i) 
+                    pixels[y, x, 0] = 255*(1+math.cos(3.32*log_iter))//2
+                    pixels[y, x, 1] = 255*(1+math.cos(0.774*log_iter))//2
+                    pixels[y, x, 2] = 255*(1+math.cos(0.412*log_iter))//2
+                    break
+                c = c * c + c0
+                i += 1
+            iters[y, x] = i                
+    return iters, pixels
+## END: numpy
+
+WIDTH = 600
+PLANE = (-2.0, -1.5, 1.0, 1.5)
+
+def validator(input_args, input_kwargs, impl_output):
+    actual_iters, actual_pixels  = impl_output
+    expected_iters, expected_pixels = mandelbrot(*input_args, **input_kwargs)
+    np.testing.assert_allclose(expected_iters, actual_iters)
+    np.testing.assert_allclose(expected_pixels, actual_pixels)
+
+def make_arrays(width):
+    iters = np.zeros((width, width), dtype=np.uint16)
+    pixels = np.zeros((width, width, 3), dtype=np.uint8)
+    return iters, pixels
+
+def input_generator():
+    for maxiter in [100, 200, 500, 1000]:
+        iters, pixels = make_arrays(WIDTH)
+        yield dict(category=("",),
+                x=maxiter,
+                input_args=(iters, pixels, PLANE, WIDTH, maxiter),
+                input_kwargs={})

examples/maths/mandelbrot_set_vectorization/bench.yaml

@@ -0,0 +1,21 @@
+name: Mandelbrot Set (ufunc)
+description: |
+    Compute the Mandelbrot Set, a set of points on the complex plane which always remain bounded by a threshold value while solving the quadratic recurrence equation. It is an iterative problem which is also compute intensive and visual in nature.  
+    In this benchmark we define a NumPy vectorized function (using Python which does not compile to C code) that takes a nested sequence of objects or numpy arrays as inputs and returns the computed output. It is compared with numba's capability of creating NumPy ufuncs without writing any C Code.  
+    Author - <a href="https://github.com/animator">Ankit Mahato</a>
+
+input_generator: numpy_vectorize.py:input_generator
+xlabel: Max iteration
+validator: numpy_vectorize.py:validator
+implementations:
+    - name: numpy_vectorize
+      description: numpy.vectorize() creates a vectorized function which uses broadcasting rules insted of for loops. Although users can write vectorized functions in Python, it is mostly for convenience and is neither optimal nor efficient.
+      function: numpy_vectorize.py:mandelbrot
+    - name: numba_vectorize
+      description: numba.vectorize() creates a NumPy func from a Python function as compared to writing C code if using the NumPy API. A func uses broadcasting rules insted of nested for loops.
+      function: numba_vectorize.py:mandelbrot
+    - name: parallel_numba_vectorize
+      description: numba.vectorize() also has the option to create a func which executes in parallel.
+      function: parallel_numba_vectorize.py:mandelbrot
+
+baseline: numpy_vectorize

examples/maths/mandelbrot_set_vectorization/numba_vectorize.py

@@ -0,0 +1,25 @@
+import numpy as np
+import numba as nb
+
+@nb.vectorize
+def get_max_iter(x, y, width, max_iter, r1, i1, r2, i2):
+    c0 = complex(r1 + (r2-r1)*x/width, 
+                 i1 + (i2-i1)*y/width)  
+    c = 0
+    for i in range(1, max_iter+1): 
+        if abs(c) > 2:  
+            break
+        c = c * c + c0
+    return i
+
+def mandelbrot(bbox, width, max_iter):  
+    grid1D = np.arange(0, width)
+    xv, yv = np.meshgrid(grid1D, grid1D)  
+    r1, i1, r2, i2 = bbox
+    iters = get_max_iter(xv, yv, width, max_iter, r1, i1, r2, i2).reshape((width, width, 1))
+    pixels = np.where(iters == np.array([max_iter]), 
+                      np.array([0, 0, 0]), 
+                      255*(1+np.cos(np.log(iters)*np.array([3.32, 0.774, 0.412])))//2)
+    pixels = pixels.astype(np.uint8)
+    iters = iters.reshape((width, width))
+    return iters, pixels