calculate_z function using @vectorize

vectorized_calc_z.py

from numba import vectorize
import numpy as np

@vectorize(['uint(c8, c8, uint)'])
def calculate_z(qi, zi, maxiter):
    output = 0.0
    for iteration in range(maxiter):
        zi = zi * zi + qi
        if abs(zi) > 2.0:
            output = iteration
            break
    return output

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 = calculate_z(q, z, maxiter)