calculate_z function using @njit

calc_z_njit.py

from numba import njit
import numpy as np

@njit
def calculate_z(q, z, output, maxiter):
    #"""Pure python with complex datatype, iterating over list of q and z"""
    for i in range(len(q)):
        zi = z[i]
        qi = q[i]
        for iteration in range(maxiter):
            zi = zi * zi + qi
            if abs(zi) > 2.0:
                output[i] = iteration
                break

def create_q_z():
    w = h = 1000

    # make a list of x and y values which will represent q
    # xx and yy are the co-ordinates, for the default configuration they'll look like:
    # if we have a 1000x1000 plot
    # xx = [-2.13, -2.1242,-2.1184000000000003, ..., 0.7526000000000064, 0.7584000000000064, 0.7642000000000064]
    # yy = [1.3, 1.2948, 1.2895999999999999, ..., -1.2844000000000058, -1.2896000000000059, -1.294800000000006]
    x1, x2, y1, y2 = -2.13, 0.77, -1.3, 1.3

    x_step = (float(x2 - x1) / float(w)) * 2
    y_step = (float(y1 - y2) / float(h)) * 2
    x = []
    y = []
    ycoord = y2
    while ycoord > y1:
        y.append(ycoord)
        ycoord += y_step
        xcoord = x1

    while xcoord < x2:
        x.append(xcoord)
        xcoord += x_step
    q = []
    for ycoord in y:
        for xcoord in x:
            q.append(complex(xcoord, ycoord))

    z = [0+0j] * len(q)
    return np.asarray(q, dtype=np.complex64), np.asarray(z, dtype=np.complex64)


maxiter = 1000
q, z = create_q_z()
output = np.zeros(len(q), dtype=np.complex64)
#calculate_z(q, z, output, maxiter)