Using numba - the code#

The implementatiom changes a little by adding the necessary bits to support correct addressing.

import math

import numpy as np
from numba import cuda

@cuda.jit
def mandel(_c, iterations, max_iterations):
    # Note: define the indexes for a 2D array and check bounds
    x, y = cuda.grid(2)
    if x < iterations.shape[0] and y < iterations.shape[1]:
        c = _c[x,y]
        c0 = c
        for iteration in range(max_iterations):
            c = c**2 + c0
            if abs(c) > 2.0:
                break
        # Note: Unusual call by reference and inplace modification
        iterations[x,y] = iteration


def calculate(x_min, x_max, y_min, y_max, max_iterations, resolution):
    x = np.linspace(x_min, x_max, resolution)
    y = np.linspace(y_min, y_max, resolution)
    c = x + y[:, None] * 1j
    iterations = np.zeros_like(c, dtype=np.uint)

    threadsperblock = (16, 16)
    blockspergrid_x = math.ceil(iterations.shape[0] / threadsperblock[0])
    blockspergrid_y = math.ceil(iterations.shape[1] / threadsperblock[1])
    blockspergrid = (blockspergrid_x, blockspergrid_y)

    mandel[blockspergrid, threadsperblock](c, iterations, max_iterations)
    return iterations, {}